Update app.py

app.py

@@ -1,25 +1,22 @@
 import os
 import gradio as gr
 import joblib
+import logfire
 import numpy as np
 import pandas as pd
 from openai import OpenAI
-from huggingface_hub import login
-from huggingface_hub import hf_hub_download
-import logfire
 from pydantic import BaseModel
 
-
+# Configure logging
 logfire.configure(token=os.getenv("LOGFIRE_API_KEY"))
 logfire.instrument_pydantic()
 
-# Load your pre-trained model and label names
-# model_path = hf_hub_download(repo_id="govtech/zoo-entry-001", filename="model.joblib", use_auth_token=True)
+# Load pre-trained model and label names
 model_data = joblib.load("model.joblib")
-model = model_data['model']
-label_names = model_data['label_names']
+model = model_data["model"]
+label_names = model_data["label_names"]
 
-class results(BaseModel):
+class Results(BaseModel):
     text: str
     hateful: float
     insults: float
@@ -28,63 +25,53 @@ class results(BaseModel):
     self_harm: float
     aom: float
 
-
 # Initialize OpenAI client
 client = OpenAI()
 
-def get_embedding(text, embedding_model="text-embedding-3-large"):
+def get_embedding(text: str, embedding_model: str = "text-embedding-3-large") -> np.ndarray:
     """
     Get embedding for the input text from OpenAI.
-    Replace newlines in the text, then call the API.
+    Replaces newlines with spaces before calling the API.
     """
     text = text.replace("\n", " ")
-    response = client.embeddings.create(
-        input=[text],
-        model=embedding_model
-    )
-    # Extract embedding vector from response
+    response = client.embeddings.create(input=[text], model=embedding_model)
     embedding = response.data[0].embedding
     return np.array(embedding)
 
-def classify_text(text):
+def classify_text(text: str):
     """
     Get the OpenAI embedding for the provided text, classify it using your model,
-    and return an updated DataFrame component with the predictions and probabilities.
+    and return a DataFrame with the rounded probabilities and binary predictions.
     """
     embedding = get_embedding(text)
-    # Add batch dimension
-    X = np.array(embedding)[None, :]
-    # Get probabilities from the model
+    X = embedding.reshape(1, -1)
     probabilities = model.predict(X)
-    # Create a DataFrame with probabilities, labels, and binary predictions
-    df = pd.DataFrame({
-        'Label': label_names,
-        'Probability': probabilities[0],
-        'Prediction': (probabilities[0] > 0.5).astype(int)
-    })
-    # Return an update to the DataFrame component to make it visible with the results
-    results(
+    rounded_probs = np.round(probabilities[0], 4)
+    
+    # Optionally log the results (this doesn't affect the output)
+    Results(
         text=text,
-        hateful=probabilities[0][0].round(4),
-        insults=probabilities[0][1].round(4),
-        sexual=probabilities[0][2].round(4),
-        violence=probabilities[0][3].round(4),
-        self_harm=probabilities[0][4].round(4),
-        aom=probabilities[0][5].round(4),
+        hateful=rounded_probs[0],
+        insults=rounded_probs[1],
+        sexual=rounded_probs[2],
+        violence=rounded_probs[3],
+        self_harm=rounded_probs[4],
+        aom=rounded_probs[5],
     )
+    
+    # Create DataFrame with rounded probabilities and binary predictions
+    df = pd.DataFrame({
+        "Label": label_names,
+        "Probability": rounded_probs,
+        "Prediction": (rounded_probs > 0.5).astype(int)
+    })
+    
     return gr.update(value=df, visible=True)
 
 with gr.Blocks(title="Zoo Entry 001") as iface:
-    
-    with gr.Row():
-        input_text = gr.Textbox(lines=5, label="Input Text")
-    
-    with gr.Row():
-        submit_btn = gr.Button("Submit")
-    
-    # Initialize the table as hidden
-    with gr.Row():
-        output_table = gr.DataFrame(label="Classification Results", visible=False)
+    input_text = gr.Textbox(lines=5, label="Input Text")
+    submit_btn = gr.Button("Submit")
+    output_table = gr.DataFrame(label="Classification Results", visible=False)
     
     submit_btn.click(fn=classify_text, inputs=input_text, outputs=output_table)