Update app.py

app.py

@@ -5,13 +5,12 @@ import gradio as gr
 from transformers import AutoTokenizer
 from huggingface_hub import hf_hub_download
 
-# Silence background warning messages
+# Silence background warning logs
 os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "1"
 
-# Target model repository
 REPO_ID = "Havoc999/distilbert-imdb-sentiment-onnx"
 
-print("Loading model and tokenizer locally...")
+print("Loading model assets locally...")
 tokenizer = AutoTokenizer.from_pretrained(REPO_ID)
 model_path = hf_hub_download(repo_id=REPO_ID, filename="model.onnx")
 session = ort.InferenceSession(model_path)
@@ -20,7 +19,7 @@ def predict_sentiment(text):
     if not text.strip():
         return {"Negative 🎬❌": 0.5, "Positive 🍿✅": 0.5}
     
-    # 1. Tokenize incoming text structure
+    # 1. Tokenize text inputs
     inputs = tokenizer(
         text, 
         return_tensors="np", 
@@ -34,26 +33,23 @@ def predict_sentiment(text):
         "attention_mask": inputs["attention_mask"].astype(np.int64)
     }
     
-    # 2. Execute local ONNX inference
-    onnx_outputs = session.run(None, onnx_inputs)
+    # 2. Force ONNX to return ONLY the classification logits array
+    onnx_outputs = session.run(["logits"], onnx_inputs)
     
-    # 3. Extract the logits array safely and flatten it to 1D
-    logits = onnx_outputs.flatten()
+    # 3. Safely convert the output to a flat 1D numpy array, regardless of nested lists
+    logits = np.asarray(onnx_outputs).flatten()
     
-    # 4. Standard Softmax conversion
+    # 4. Softmax probability conversion
     exp_logits = np.exp(logits - np.max(logits))
     probabilities = exp_logits / np.sum(exp_logits)
     
-    # 5. Extract items individually to guarantee conversion to pure Python floats
-    neg_score = float(probabilities)
-    pos_score = float(probabilities)
-    
+    # 5. Return pure Python float scalars to the Gradio UI
     return {
-        "Negative 🎬❌": neg_score,
-        "Positive 🍿✅": pos_score
+        "Negative 🎬❌": float(probabilities),
+        "Positive 🍿✅": float(probabilities)
     }
 
-# 6. Minimal Gradio layout configuration
+# 6. Initialize UI
 demo = gr.Interface(
     fn=predict_sentiment,
     inputs=gr.Textbox(lines=4, placeholder="Type an IMDB movie review here...", label="Your Movie Review"),