checkpointing changes

app.py

@@ -11,6 +11,8 @@ Version Log:
     - Read sysprompt and initial_message from file
     - Begins with user entering name/alias
     - Azure OpenAI?
+- 2024.01.31: wandb does not work for use case, what to do instead?
+    - Write to local file and then upload at end? (does filestream cause blocking?)
 
 """
 from __future__ import annotations
@@ -23,7 +25,6 @@ import gradio as gr
 import openai
 
 from pathlib import Path
-from base64 import urlsafe_b64decode
 
 logger = logging.getLogger(__name__)
 
@@ -31,7 +32,8 @@ from utils import PromptTemplate, convert_gradio_to_openai, seed_openai_key
 
 
 # %% Initialization
-CONFIG_DIR: Path = Path("./SPIA2024")
+
+CONFIG_DIR: Path = Path("./SPIA2024_localtesting")
 if os.environ.get(f"OPENAI_API_KEY", "DEFAULT") == "DEFAULT":
     seed_openai_key()
 client = openai.OpenAI()
@@ -53,9 +55,15 @@ def load_config(
 
 
 def initialize_interview(
-    system_message: str, initial_message: str, model_args: dict[str, str | float]
-) -> tuple[gr.Chatbot, gr.Textbox, gr.Button, gr.Button]:
-    "Read system prompt and start interview"
+    system_message: str,
+    initial_message: str,
+    model_args: dict[str, str | float]
+) -> tuple[gr.Chatbot,
+           gr.Textbox,
+           gr.Button,
+           gr.Button,
+           gr.Button]:
+    "Read system prompt and start interview. Change visibilities of elements."
     if len(initial_message) == 0:  # If empty inital message, ask the LM to write it.
         initial_message = call_openai(
             [], system_message, client, model_args, stream=False
@@ -64,7 +72,7 @@ def initialize_interview(
         [None, initial_message]
     ]  # First item is for user, in this case bot starts interaction.
     return (
-        gr.Chatbot(visible=True, value=chat_history),
+        gr.Chatbot(visible=True, value=chat_history), # chatDisplay
         gr.Textbox(
             placeholder="Type response here. Hit 'enter' to submit.",
             visible=True,
@@ -72,7 +80,6 @@ def initialize_interview(
         ),  # chatInput
         gr.Button(visible=True, interactive=True),  # chatSubmit
         gr.Button(visible=False),  # startInterview
-        gr.Textbox(visible=False),  # userBox
         gr.Button(visible=True),  # resetButton
     )
 
@@ -82,18 +89,20 @@ def initialize_tracker(
     system_message: PromptTemplate,
     userid: str,
     wandb_args: dict[str, str | list[str]],
-) -> None:
-    "Initializes wandb run for interview"
+) -> gr.Textbox:
+    "Initializes wandb run for interview. Resets userBox afterwards."
     run_config = model_args | {
         "system_message": str(system_message),
         "userid": userid,
     }
+    logger.info(f"Initializing WandB run for {userid}")
     wandb.init(
         project=wandb_args["project"],
         name=userid,
         config=run_config,
         tags=wandb_args["tags"],
     )
+    return gr.Textbox(value=None, visible=False)
 
 
 def save_interview(
@@ -197,6 +206,7 @@ def reset_interview() -> (
 # LAYOUT
 with gr.Blocks() as demo:
     gr.Markdown("# StewartLab LM Interviewer")
+    userDisplay = gr.Markdown("")
 
     # Config values
     configDir = gr.State(value=CONFIG_DIR)
@@ -207,7 +217,7 @@ with gr.Blocks() as demo:
 
     ## Start interview by entering name or alias
     userBox = gr.Textbox(
-        placeholder="Enter name or alias and hit 'enter' to begin.", show_label=False
+        value=None, placeholder="Enter name or alias and hit 'enter' to begin.", show_label=False
     )
     startInterview = gr.Button("Start Interview", variant="primary", visible=True)
 
@@ -232,6 +242,7 @@ with gr.Blocks() as demo:
 
     ## INTERACTIONS
     # Start Interview button
+    userBox.change(lambda x: x, inputs=[userBox], outputs=[userDisplay], show_progress=False)
     userBox.submit(
         load_config,
         inputs=configDir,
@@ -244,12 +255,14 @@ with gr.Blocks() as demo:
             chatInput,
             chatSubmit,
             startInterview,
-            userBox,
             resetButton,
         ],
     ).then(
-        initialize_tracker, inputs=[modelArgs, systemMessage, userBox, wandbArgs]
+        initialize_tracker,
+        inputs=[modelArgs, systemMessage, userBox, wandbArgs],
+        outputs=[userBox]
     )
+
     startInterview.click(
         load_config,
         inputs=configDir,
@@ -257,9 +270,17 @@ with gr.Blocks() as demo:
     ).then(
         initialize_interview,
         inputs=[systemMessage, initialMessage, modelArgs],
-        outputs=[chatDisplay, chatInput, chatSubmit, startInterview, resetButton],
+        outputs=[
+            chatDisplay,
+            chatInput,
+            chatSubmit,
+            startInterview,
+            resetButton,
+        ],
     ).then(
-        initialize_tracker, inputs=[modelArgs, systemMessage, userBox, wandbArgs]
+        initialize_tracker,
+        inputs=[modelArgs, systemMessage, userBox, wandbArgs],
+        outputs=[userBox]
     )
 
     # Chat interaction

utils.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import os
+import openai
 from configparser import ConfigParser
 from pathlib import Path
 from string import Formatter
@@ -105,3 +106,34 @@ def seed_openai_key(cfg: str = "~/.cfg/openai.cfg") -> None:
     except:
         raise ValueError(f"Could not using read file at: {cfg}.")
     os.environ["OPENAI_API_KEY"] = config["API_KEY"]["secret"]
+
+
+
+def query_openai(
+    messages: list[dict[str, str]],
+    system_message: str | None,
+    client: openai.Client,
+    model_args: dict,
+    stream: bool = False,
+) -> None:
+    "Utility function for calling OpenAI chat. Expects formatted messages."
+    if not messages:
+        messages = []
+    if system_message:
+        messages = [{"role": "system", "content": system_message}] + messages
+    try:
+        response = client.chat.completions.create(
+            messages=messages, **model_args, stream=stream
+        )
+        if stream:
+            for chunk in response:
+                yield chunk.choices[0].message.content
+        else:
+            content = response.choices[0].message.content
+            return content
+    except openai.APIConnectionError | openai.APIStatusError as e:
+        error_msg = (
+            "API unreachable.\n" f"STATUS_CODE: {e.status_code}" f"ERROR: {e.response}"
+        )
+    except openai.RateLimitError:
+        warning_msg = "Hit rate limit. Wait a moment and retry."