Update app.py

app.py

@@ -8,116 +8,109 @@ import gradio as gr
 
 class ONNXInferencePipeline:
     def __init__(self, repo_id):
-        # Download files from Hugging Face Hub using the provided repo_id.
-        # Note: The ONNX model file is now "moodmeter.onnx"
-        self.onnx_path = hf_hub_download(repo_id=repo_id, filename="minddmeter.onnx")
+        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")
 
-        # Load configuration from the hyperparameters file.
-        # Ensure that the JSON contains a "tokenizer" key with "max_len" defined.
+        # Load configuration from hyperparameters.json.
+        # This file should have a key "MAX_LEN" specifying the maximum token sequence length.
         with open(self.config_path, "r") as f:
             self.config = json.load(f)
 
-        # Initialize the tokenizer using the downloaded file.
+        # Initialize the tokenizer from file.
         self.tokenizer = Tokenizer.from_file(self.tokenizer_path)
-        self.max_len = self.config["tokenizer"]["max_len"]
+        # Use the maximum sequence length from the hyperparameters.
+        self.max_len = self.config["MAX_LEN"]
 
-        # Initialize the ONNX runtime session with the downloaded model.
+        # Initialize the ONNX runtime session.
         self.session = ort.InferenceSession(self.onnx_path)
-        # Use CUDA if available; otherwise, default to CPU.
+        # Use CUDA if available, otherwise default to CPU.
         self.providers = ['CPUExecutionProvider']
         if 'CUDAExecutionProvider' in ort.get_available_providers():
             self.providers = ['CUDAExecutionProvider']
         self.session.set_providers(self.providers)
 
     def preprocess(self, text):
-        """
-        Tokenize the input text, truncate or pad it to the maximum length, and return a numpy array.
-        """
         encoding = self.tokenizer.encode(text)
-        # Truncate token ids to self.max_len tokens.
+        # Truncate to self.max_len tokens
         ids = encoding.ids[:self.max_len]
-        # Pad with zeros if the sequence is shorter than self.max_len.
+        # 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:
-          - The predicted label (mapped from the model output),
-          - The confidence of that prediction,
-          - The full list of probabilities.
+        Given an input text string, run inference and return only the granular stress level.
+        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"
         """
-        # Preprocess the text
         input_array = self.preprocess(text)
 
-        # Run inference using the ONNX runtime session.
-        results = self.session.run(None, {"input": input_array})
+        outputs = self.session.run(None, {"input": input_array})
+        logits = outputs[0]
 
-        # Compute softmax probabilities from the logits.
-        logits = results[0]
         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 the actual class label.
-        # Here we assume the model outputs: 0 -> "neg" and 1 -> "pos"
-        label_mapping = {'neg': 'Negative', 'pos': 'Positive'}
-        class_labels = ['neg', 'pos']
-        predicted_label = label_mapping[class_labels[predicted_class]]
+        class_labels = ["Not Stressed", "Stressed"]
+        predicted_label = class_labels[predicted_class]
         confidence = float(probabilities[0][predicted_class])
 
-        return {
-            'label': predicted_label,
-            'confidence': confidence,
-            'probabilities': probabilities[0].tolist()
-        }
+        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.65:
+                stress_level = "Low Stress"
+            elif 0.65 <= stress_confidence < 0.90:
+                stress_level = "Moderate Stress"
+            else:  # 0.90 ≤ stress_confidence ≤ 1.00
+                stress_level = "High Stress"
+        else:
+            stress_level = "Not Stressed"
+
+        return {"stress_level": stress_level}
 
 
-# Example usage and Gradio Interface
 if __name__ == "__main__":
-    # Initialize the pipeline with the correct Hugging Face repository ID.
     pipeline = ONNXInferencePipeline(repo_id="iimran/MindMeter")
 
-    # Example texts for testing (stress/sentiment related)
-    example_texts = [
-        "I feel the constant strain of financial hardship, and the stress of it all leaves me emotionally drained.",
-        "I am stressed by the thought of unexpected expenses derailing my already fragile budget. Can I request council to make an installment plan for me?",
-        "The burden of my debt is immense, and the stress from my financial hardships is unbearable. Can I request council to waive my rates fee?"
-    ]
-    
-    for text in example_texts:
-        result = pipeline.predict(text)
-        print(f"Input: {text}")
-        print(f"Prediction: {result['label']} ({result['confidence']:.2%})")
-        print(f"Confidence Scores: Negative={result['probabilities'][0]:.2%}, Positive={result['probabilities'][1]:.2%}")
-        print("-" * 80)
+    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 and i wanna end it"
+
+    result1 = pipeline.predict(text1)
+    result2 = pipeline.predict(text2)
+    result3 = pipeline.predict(text3)
+
+    print(f"Prediction for text 1: {result1}")
+    print(f"Prediction for text 2: {result2}")
+    print(f"Prediction for text 3: {result3}")
 
-    # Define a function for Gradio to call.
     def gradio_predict(text):
         result = pipeline.predict(text)
-        return (
-            f"Prediction: {result['label']} ({result['confidence']:.2%})\n"
-            f"Confidence Scores: Negative={result['probabilities'][0]:.2%}, Positive={result['probabilities'][1]:.2%}"
-        )
+        return result["stress_level"]
 
-    # Create the Gradio interface.
     iface = gr.Interface(
         fn=gradio_predict,
-        inputs=gr.Textbox(lines=7, placeholder="Enter text here..."),
+        inputs=gr.Textbox(lines=7, placeholder="Enter your text here..."),
         outputs="text",
-        title="MindMeter – Service-Focused Stress Identification Agent",
+        title="MindMeter – Granular Stress Level Predictor",
         description=(
-            "MindMeter evaluates the mental and emotional tone of your text input. "
-            "It analyzes the content to provide an indication of the overall stress element in the input, "
-            "displaying a Low Stressed, Medium Stressed, High Stressed or not stressed scores. "
-            "Enter your text below to see the analysis."
+            "MindMeter predicts the granular stress level from your text input. "
+            "The possible outputs are: Not Stressed, Low Stress, Moderate Stress, or High Stress."
         ),
         examples=[
-            "I feel the constant strain of financial hardship, and the stress of it all leaves me emotionally drained.",
-            "I am stressed by the thought of unexpected expenses derailing my already fragile budget. Can I request council to make an installment plan for me?",
-            "The burden of my debt is immense, and the stress from my financial hardships is unbearable. Can I request council to waive my rates fee?"
+            "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"
         ]
     )