app.py

import json
import numpy as np
from tokenizers import Tokenizer
import onnxruntime as ort
from huggingface_hub import hf_hub_download
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.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.
        with open(self.config_path, "r") as f:
            self.config = json.load(f)

        # Initialize the tokenizer using the downloaded file.
        self.tokenizer = Tokenizer.from_file(self.tokenizer_path)
        self.max_len = self.config["tokenizer"]["max_len"]

        # Initialize the ONNX runtime session with the downloaded model.
        self.session = ort.InferenceSession(self.onnx_path)
        # 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.
        ids = encoding.ids[:self.max_len]
        # Pad with zeros if the sequence is shorter than self.max_len.
        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.
        """
        # Preprocess the text
        input_array = self.preprocess(text)

        # Run inference using the ONNX runtime session.
        results = self.session.run(None, {"input": input_array})

        # 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]]
        confidence = float(probabilities[0][predicted_class])

        return {
            'label': predicted_label,
            'confidence': confidence,
            'probabilities': probabilities[0].tolist()
        }


# 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)

    # 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%}"
        )

    # Create the Gradio interface.
    iface = gr.Interface(
        fn=gradio_predict,
        inputs=gr.Textbox(lines=7, placeholder="Enter text here..."),
        outputs="text",
        title="MindMeter – Service-Focused Stress Identification Agent",
        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."
        ),
        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?"
        ]
    )

    # Launch the Gradio app.
    iface.launch()