import os
import numpy as np
import onnxruntime as ort
import gradio as gr
from transformers import AutoTokenizer
from huggingface_hub import hf_hub_download

# Silence background warning logs
os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "1"

REPO_ID = "Havoc999/distilbert-imdb-sentiment-onnx"

# Load local model assets
tokenizer = AutoTokenizer.from_pretrained(REPO_ID)
model_path = hf_hub_download(repo_id=REPO_ID, filename="model.onnx")
session = ort.InferenceSession(model_path)

def predict_sentiment(text):
    if not text.strip():
        return {"Negative 🎬❌": 0.5, "Positive 🍿✅": 0.5}
    
    # 1. Tokenize text
    inputs = tokenizer(
        text, 
        return_tensors="np", 
        padding="max_length", 
        truncation=True, 
        max_length=512
    )
    
    onnx_inputs = {
        "input_ids": inputs["input_ids"].astype(np.int64),
        "attention_mask": inputs["attention_mask"].astype(np.int64)
    }
    
    # 2. Run local ONNX inference
    onnx_outputs = session.run(None, onnx_inputs)
    logits = np.asarray(onnx_outputs).flatten()
    
    # 3. Softmax math conversion
    exp_logits = np.exp(logits - np.max(logits))
    scores_array = exp_logits / np.sum(exp_logits)
    
    # 4. Convert NumPy array directly to standard Python floats inside a list
    clean_scores = scores_array.tolist()
    
    # 5. NO INDEXES OR FLICKERING TYPOS: Map labels directly to values
    labels = ["Negative 🎬❌", "Positive 🍿✅"]
    return dict(zip(labels, clean_scores))

# 6. Set up the web interface
demo = gr.Interface(
    fn=predict_sentiment,
    inputs=gr.Textbox(lines=4, placeholder="Will it hurt Someone? type to check if it is a positive/ negative comment", label="Your comment"),
    outputs=gr.Label(num_top_classes=2, label="Sentiment Predictions"),
    title="🎬 IMDB Movie Review Sentiment Classifier",
    description="Optimized local CPU inference using ONNX Runtime.",
    theme="soft"
)

if __name__ == "__main__":
    demo.launch()