Update app.py

app.py

@@ -7,11 +7,11 @@ from openai import OpenAI
 from huggingface_hub import login
 from huggingface_hub import hf_hub_download
 import logfire
+from pydantic import BaseModel
 
 
 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)
@@ -19,6 +19,16 @@ model_data = joblib.load("model.joblib")
 model = model_data['model']
 label_names = model_data['label_names']
 
+class results(BaseModel):
+    text: str
+    hateful: float
+    insults: float
+    sexual: float
+    violence: float
+    self_harm: float
+    aom: float
+
+
 # Initialize OpenAI client
 client = OpenAI()
 
@@ -53,7 +63,15 @@ def classify_text(text):
         'Prediction': (probabilities[0] > 0.5).astype(int)
     })
     # Return an update to the DataFrame component to make it visible with the results
-    logfire.info(f"{text} ({probabilities[0]})")
+    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),
+    )
     return gr.update(value=df, visible=True)
 
 with gr.Blocks(title="Zoo Entry 001") as iface: