Upload 12 files

openai_fake_data_generator.py

@@ -2,54 +2,45 @@ from collections import namedtuple
 from typing import Optional
 
 import openai
+from openai import OpenAI, AzureOpenAI
 import logging
 
 logger = logging.getLogger("presidio-streamlit")
 
 OpenAIParams = namedtuple(
     "open_ai_params",
-    ["openai_key", "model", "api_base", "deployment_name", "api_version", "api_type"],
+    ["openai_key", "model", "api_base", "deployment_id", "api_version", "api_type"],
 )
 
 
-def set_openai_params(openai_params: OpenAIParams):
-    """Set the OpenAI API key.
-    :param openai_params: OpenAIParams object with the following fields: key, model, api version, deployment_name,
-    The latter only relate to Azure OpenAI deployments.
-    """
-    openai.api_key = openai_params.openai_key
-    openai.api_version = openai_params.api_version
-    if openai_params.api_base:
-        openai.api_base = openai_params.api_base
-        openai.api_type = openai_params.api_type
-
-
 def call_completion_model(
     prompt: str,
-    model: str = "text-davinci-003",
-    max_tokens: int = 512,
-    deployment_id: Optional[str] = None,
+    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 model: OpenAI model name
-    :param max_tokens: Model's max_tokens parameter
-    :param deployment_id: Azure OpenAI deployment ID
+    :param openai_params: OpenAI parameters for the completion model
+    :param max_tokens: The maximum number of tokens to generate.
     """
-    if deployment_id:
-        response = openai.Completion.create(
-            deployment_id=deployment_id,
-            model=model,
-            prompt=prompt,
-            max_tokens=max_tokens,
+    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,
         )
     else:
-        response = openai.Completion.create(
-            model=model, prompt=prompt, max_tokens=max_tokens
-        )
+        client = OpenAI(api_key=openai_params.openai_key)
+
+    response = client.completions.create(
+        model=openai_params.model,
+        prompt=prompt,
+        max_tokens=max_tokens,
+    )
 
-    return response["choices"][0].text
+    return response.choices[0].text.strip()
 
 
 def create_prompt(anonymized_text: str) -> str:

presidio_helpers.py

@@ -16,10 +16,9 @@ from presidio_anonymizer import AnonymizerEngine
 from presidio_anonymizer.entities import OperatorConfig
 
 from openai_fake_data_generator import (
-    set_openai_params,
     call_completion_model,
-    create_prompt,
     OpenAIParams,
+    create_prompt,
 )
 from presidio_nlp_engine_config import (
     create_nlp_engine_with_spacy,
@@ -218,14 +217,9 @@ def create_fake_data(
     if not openai_params.openai_key:
         return "Please provide your OpenAI key"
     results = anonymize(text=text, operator="replace", analyze_results=analyze_results)
-    set_openai_params(openai_params)
     prompt = create_prompt(results.text)
     print(f"Prompt: {prompt}")
-    fake = call_openai_api(
-        prompt=prompt,
-        openai_model_name=openai_params.model,
-        openai_deployment_name=openai_params.deployment_name,
-    )
+    fake = call_completion_model(prompt=prompt, openai_params=openai_params)
     return fake
 
 

presidio_streamlit.py

@@ -135,7 +135,8 @@ def set_up_openai_synthesis():
             "Azure OpenAI base URL",
             value=os.getenv("AZURE_OPENAI_ENDPOINT", default=""),
         )
-        st_deployment_name = st.sidebar.text_input(
+        openai_key = os.getenv("AZURE_OPENAI_KEY", default="")
+        st_deployment_id = st.sidebar.text_input(
             "Deployment name", value=os.getenv("AZURE_OPENAI_DEPLOYMENT", default="")
         )
         st_openai_version = st.sidebar.text_input(
@@ -143,11 +144,13 @@ def set_up_openai_synthesis():
             value=os.getenv("OPENAI_API_VERSION", default="2023-05-15"),
         )
     else:
-        st_openai_version = openai_api_type = st_openai_api_base = None
-        st_deployment_name = ""
+        openai_api_type = "openai"
+        st_openai_version = st_openai_api_base = None
+        st_deployment_id = ""
+        openai_key = os.getenv("OPENAI_KEY", default="")
     st_openai_key = st.sidebar.text_input(
         "OPENAI_KEY",
-        value=os.getenv("OPENAI_KEY", default=""),
+        value=openai_key,
         help="See https://help.openai.com/en/articles/4936850-where-do-i-find-my-secret-api-key for more info.",
         type="password",
     )
@@ -159,7 +162,7 @@ def set_up_openai_synthesis():
     return (
         openai_api_type,
         st_openai_api_base,
-        st_deployment_name,
+        st_deployment_id,
         st_openai_version,
         st_openai_key,
         st_openai_model,
@@ -179,7 +182,7 @@ elif st_operator == "synthesize":
     (
         openai_api_type,
         st_openai_api_base,
-        st_deployment_name,
+        st_deployment_id,
         st_openai_version,
         st_openai_key,
         st_openai_model,
@@ -189,7 +192,7 @@ elif st_operator == "synthesize":
         openai_key=st_openai_key,
         model=st_openai_model,
         api_base=st_openai_api_base,
-        deployment_name=st_deployment_name,
+        deployment_id=st_deployment_id,
         api_version=st_openai_version,
         api_type=openai_api_type,
     )