Update app.py

app.py

@@ -8,6 +8,7 @@ import gradio as gr
 
 class ONNXInferencePipeline:
     def __init__(self, repo_id):
+        # Download the model files from Hugging Face Hub.
         self.onnx_path = hf_hub_download(repo_id=repo_id, filename="mindmeter.onnx")
         self.tokenizer_path = hf_hub_download(repo_id=repo_id, filename="train_bpe_tokenizer.json")
         self.config_path = hf_hub_download(repo_id=repo_id, filename="hyperparameters.json")
@@ -19,7 +20,7 @@ class ONNXInferencePipeline:
 
         # Initialize the tokenizer from file.
         self.tokenizer = Tokenizer.from_file(self.tokenizer_path)
-        # Use the maximum sequence length from the hyperparameters.
+        # Use the maximum sequence length from the configuration.
         self.max_len = self.config["MAX_LEN"]
 
         # Initialize the ONNX runtime session.
@@ -31,57 +32,46 @@ class ONNXInferencePipeline:
         self.session.set_providers(self.providers)
 
     def preprocess(self, text):
+        """
+        Tokenize the input text, truncate or pad it to max_len, and return a numpy array.
+        """
         encoding = self.tokenizer.encode(text)
-        # Truncate to self.max_len tokens
+        # Truncate to self.max_len tokens.
         ids = encoding.ids[:self.max_len]
-        # Pad with zeros if necessary
+        # Pad with zeros if necessary.
         padding = [0] * (self.max_len - len(ids))
         return np.array(ids + padding, dtype=np.int64).reshape(1, -1)
 
     def predict(self, text):
         """
-        Given an input text string, run inference and return only the granular stress level.
+        Given an input text string, run inference and return the predicted stress label.
         The model outputs:
           0 -> "Not Stressed"
           1 -> "Stressed"
-        When the model predicts "Stressed", a confidence-based thresholding is applied:
-          - confidence < 0.40: "Not Stressed" (fallback)
-          - 0.40 ≤ confidence < 0.65: "Low Stress"
-          - 0.65 ≤ confidence < 0.90: "Moderate Stress"
-          - 0.90 ≤ confidence: "High Stress"
+        This function returns one of these two labels.
         """
         input_array = self.preprocess(text)
-
         outputs = self.session.run(None, {"input": input_array})
         logits = outputs[0]
 
+        # Compute softmax probabilities from the logits.
         exp_logits = np.exp(logits)
         probabilities = exp_logits / np.sum(exp_logits, axis=1, keepdims=True)
         predicted_class = int(np.argmax(probabilities))
+        
+        # Map the predicted index to a label.
         class_labels = ["Not Stressed", "Stressed"]
         predicted_label = class_labels[predicted_class]
-        confidence = float(probabilities[0][predicted_class])
-
-        if predicted_label == "Stressed":
-            # Use the confidence of the "Stressed" class
-            stress_confidence = confidence
-            if stress_confidence < 0.40:
-                stress_level = "Not Stressed"  # Fallback (unlikely)
-            elif 0.40 <= stress_confidence < 0.60:
-                stress_level = "Low Stress"
-            elif 0.60 <= stress_confidence < 0.95:
-                stress_level = "Moderate Stress"
-            else:  # 0.95 ≤ stress_confidence ≤ 1.00
-                stress_level = "High Stress"
-        else:
-            stress_level = "Not Stressed"
 
-        return {"stress_level": stress_level}
+        # Return only the predicted stress label.
+        return {"stress_level": predicted_label}
 
 
 if __name__ == "__main__":
+    # Initialize the pipeline with the Hugging Face repository ID.
     pipeline = ONNXInferencePipeline(repo_id="iimran/MindMeter")
 
+    # Example input texts for local testing.
     text1 = "Yay! what a happy life"
     text2 = "I’ve missed another loan payment, and I don’t know how I’m going to catch up. The pressure is unbearable."
     text3 = "I am upset about how badly life is trating me these days, its shit."
@@ -94,23 +84,25 @@ if __name__ == "__main__":
     print(f"Prediction for text 2: {result2}")
     print(f"Prediction for text 3: {result3}")
 
+    # Define a function for Gradio to use.
     def gradio_predict(text):
         result = pipeline.predict(text)
         return result["stress_level"]
 
+    # Create the Gradio interface.
     iface = gr.Interface(
         fn=gradio_predict,
         inputs=gr.Textbox(lines=7, placeholder="Enter your text here..."),
         outputs="text",
-        title="MindMeter – Granular Stress Level Detection Agent",
+        title="MindMeter – Stress Detection Agent",
         description=(
-            "MindMeter is designed to swiftly assess the stress levels expressed in text communications, making it an invaluable tool for local councils, especially when addressing financial hardship cases. By analyzing the tone and wording in emails or chat messages, MindMeter categorizes the expressed sentiment into one of four levels—Not Stressed, Low Stress, Moderate Stress, or High Stress. This allows council representatives to quickly identify residents who might be under significant stress due to financial challenges. In turn, councils can prioritize outreach and tailor support services to address urgent concerns effectively, ensuring that vulnerable community members receive the timely assistance they need."
+            "MindMeter is designed to swiftly assess the stress levels expressed in text communications, making it an invaluable tool for local councils, especially when addressing financial hardship cases. By analyzing the tone and wording in emails or chat messages, MindMeter categorizes the expressed sentiment into  'Stressed' or 'Not Stressed' outputs. This allows council representatives to quickly identify residents who might be under significant stress due to financial challenges. In turn, councils can prioritize outreach and tailor support services to address urgent concerns effectively, ensuring that vulnerable community members receive the timely assistance they need."
             " This agent will identify Risk from Text communication - Next Agent will read the financial reports/bank statements/loan statements and cross verify the finacnail hardship."
         ),
         examples=[
             "Yay! what a happy life",
             "I’ve missed another loan payment, and I don’t know how I’m going to catch up. The pressure is unbearable.",
-            "I am upset about how badly life is trating me these days, its shit and i wanna end it"
+            "I am upset about how badly life is trating me these days, its shit."
         ]
     )