Update openai_fake_data_generator.py

openai_fake_data_generator.py

@@ -1,8 +1,6 @@
 from collections import namedtuple
 from typing import Optional
-
 import openai
-from openai import OpenAI, AzureOpenAI
 import logging
 
 logger = logging.getLogger("presidio-streamlit")
@@ -12,50 +10,42 @@ OpenAIParams = namedtuple(
     ["openai_key", "model", "api_base", "deployment_id", "api_version", "api_type"],
 )
 
-
 def call_completion_model(
     prompt: str,
     openai_params: OpenAIParams,
     max_tokens: Optional[int] = 256,
 ) -> str:
-    """Creates a request for the OpenAI Completion service and returns the response.
-
-    :param prompt: The prompt for the completion model
-    :param openai_params: OpenAI parameters for the completion model
-    :param max_tokens: The maximum number of tokens to generate.
-    """
+    """Creates a request for the OpenAI Completion service and returns the response."""
+    openai.api_key = openai_params.openai_key
+    
     if openai_params.api_type.lower() == "azure":
-        client = AzureOpenAI(
-            api_version=openai_params.api_version,
-            api_key=openai_params.openai_key,
-            azure_endpoint=openai_params.api_base,
-            azure_deployment=openai_params.deployment_id,
+        openai.api_base = openai_params.api_base
+        openai.api_version = openai_params.api_version
+        openai.api_type = "azure"
+        response = openai.Completion.create(
+            engine=openai_params.deployment_id,
+            prompt=prompt,
+            max_tokens=max_tokens,
         )
     else:
-        client = OpenAI(api_key=openai_params.openai_key)
-
-    response = client.completions.create(
-        model=openai_params.model,
-        prompt=prompt,
-        max_tokens=max_tokens,
-    )
+        response = openai.Completion.create(
+            model=openai_params.model,
+            prompt=prompt,
+            max_tokens=max_tokens,
+        )
 
     return response.choices[0].text.strip()
 
-
 def create_prompt(anonymized_text: str) -> str:
     """
     Create the prompt with instructions to GPT-3.
-
-    :param anonymized_text: Text with placeholders instead of PII values, e.g. My name is <PERSON>.
+    :param anonymized_text: Text with placeholders instead of PII values.
     """
-
     prompt = f"""
     Your role is to create synthetic text based on de-identified text with placeholders instead of Personally Identifiable Information (PII).
     Replace the placeholders (e.g. ,<PERSON>, {{DATE}}, {{ip_address}}) with fake values.
 
     Instructions:
-
     a. Use completely random numbers, so every digit is drawn between 0 and 9.
     b. Use realistic names that come from diverse genders, ethnicities and countries.
     c. If there are no placeholders, return the text as is.
@@ -63,13 +53,6 @@ def create_prompt(anonymized_text: str) -> str:
     e. If PII exists in the input, replace it with fake values in the output.
     f. Remove whitespace before and after the generated text
     
-    input: [[TEXT STARTS]] How do I change the limit on my credit card {{credit_card_number}}?[[TEXT ENDS]]
-    output: How do I change the limit on my credit card 2539 3519 2345 1555?
-    input: [[TEXT STARTS]]<PERSON> was the chief science officer at <ORGANIZATION>.[[TEXT ENDS]]
-    output: Katherine Buckjov was the chief science officer at NASA.
-    input: [[TEXT STARTS]]Cameroon lives in <LOCATION>.[[TEXT ENDS]]
-    output: Vladimir lives in Moscow.
-    
     input: [[TEXT STARTS]]{anonymized_text}[[TEXT ENDS]]
     output:"""
     return prompt