Update app.py

app.py

@@ -3,66 +3,80 @@ import numpy as np
 import onnxruntime as ort
 from PIL import Image
 from huggingface_hub import hf_hub_download
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, AutoConfig
+
+
+# ---------------------------------------------------------
+# Helper
+# ---------------------------------------------------------
+def softmax(x):
+    e = np.exp(x - np.max(x))
+    return e / e.sum()
+
 
 # ---------------------------------------------------------
-# Download ONNX files directly from Hugging Face model hubs
+# Load ONNX models + tokenizers + configs (for labels)
 # ---------------------------------------------------------
 
-# 1) Multilingual sentiment (DistilBERT)
-multilingual_onnx = hf_hub_download(
+# --- Model 1: Multilingual DistilBERT Sentiment ---
+multilingual_onnx_path = hf_hub_download(
     repo_id="lxyuan/distilbert-base-multilingual-cased-sentiments-student",
     filename="onnx/model.onnx"
 )
 tokenizer_multilingual = AutoTokenizer.from_pretrained(
     "lxyuan/distilbert-base-multilingual-cased-sentiments-student"
 )
-session_multilingual = ort.InferenceSession(multilingual_onnx, providers=["CPUExecutionProvider"])
+config_multilingual = AutoConfig.from_pretrained(
+    "lxyuan/distilbert-base-multilingual-cased-sentiments-student"
+)
+labels_multilingual = config_multilingual.id2label
+session_multilingual = ort.InferenceSession(multilingual_onnx_path, providers=["CPUExecutionProvider"])
 
-# 2) SDG BERT
-sdgbert_onnx = hf_hub_download(
+
+# --- Model 2: SDG-BERT ---
+sdg_onnx_path = hf_hub_download(
     repo_id="sadickam/sdgBERT",
     filename="onnx/model.onnx"
 )
 tokenizer_sdg = AutoTokenizer.from_pretrained("sadickam/sdgBERT")
-session_sdg = ort.InferenceSession(sdgbert_onnx, providers=["CPUExecutionProvider"])
+config_sdg = AutoConfig.from_pretrained("sadickam/sdgBERT")
+labels_sdg = config_sdg.id2label
+session_sdg = ort.InferenceSession(sdg_onnx_path, providers=["CPUExecutionProvider"])
+
 
-# 3) German sentiment BERT
-german_onnx = hf_hub_download(
+# --- Model 3: German Sentiment BERT ---
+german_onnx_path = hf_hub_download(
     repo_id="oliverguhr/german-sentiment-bert",
     filename="onnx/model.onnx"
 )
 tokenizer_german = AutoTokenizer.from_pretrained("oliverguhr/german-sentiment-bert")
-session_german = ort.InferenceSession(german_onnx, providers=["CPUExecutionProvider"])
+config_german = AutoConfig.from_pretrained("oliverguhr/german-sentiment-bert")
+labels_german = config_german.id2label
+session_german = ort.InferenceSession(german_onnx_path, providers=["CPUExecutionProvider"])
+
 
-# 4) ViT-small image classifier
-vit_onnx = hf_hub_download(
+# --- Model 4: ViT Image Classifier ---
+vit_onnx_path = hf_hub_download(
     repo_id="WinKawaks/vit-small-patch16-224",
     filename="onnx/model.onnx"
 )
-session_vit = ort.InferenceSession(vit_onnx, providers=["CPUExecutionProvider"])
+session_vit = ort.InferenceSession(vit_onnx_path, providers=["CPUExecutionProvider"])
 
-# Basic preprocessing params (ImageNet)
 IMAGE_SIZE = 224
 MEAN = [0.485, 0.456, 0.406]
-STD  = [0.229, 0.224, 0.225]
+STD = [0.229, 0.224, 0.225]
 
 
 # ---------------------------------------------------------
 # Inference functions
 # ---------------------------------------------------------
-def softmax(x):
-    e = np.exp(x - np.max(x))
-    return e / e.sum(axis=-1, keepdims=True)
-
 
 def run_multilingual(text):
     inputs = tokenizer_multilingual(text, return_tensors="np", truncation=True, padding=True)
     inputs = {k: v.astype(np.int64) for k, v in inputs.items()}
     logits = session_multilingual.run(None, inputs)[0][0]
     probs = softmax(logits)
-    labels = tokenizer_multilingual.model.config.id2label
-    return {labels[i]: float(probs[i]) for i in range(len(probs))}
+    return {labels_multilingual[i]: float(probs[i]) for i in range(len(probs))}
 
 
 def run_sdg(text):
@@ -70,8 +84,7 @@ def run_sdg(text):
     inputs = {k: v.astype(np.int64) for k, v in inputs.items()}
     logits = session_sdg.run(None, inputs)[0][0]
     probs = softmax(logits)
-    labels = tokenizer_sdg.model.config.id2label
-    return {labels[i]: float(probs[i]) for i in range(len(probs))}
+    return {labels_sdg[i]: float(probs[i]) for i in range(len(probs))}
 
 
 def run_german(text):
@@ -79,8 +92,7 @@ def run_german(text):
     inputs = {k: v.astype(np.int64) for k, v in inputs.items()}
     logits = session_german.run(None, inputs)[0][0]
     probs = softmax(logits)
-    labels = tokenizer_german.model.config.id2label
-    return {labels[i]: float(probs[i]) for i in range(len(probs))}
+    return {labels_german[i]: float(probs[i]) for i in range(len(probs))}
 
 
 def preprocess_vit(image):
@@ -101,35 +113,52 @@ def run_vit(image):
 
 
 # ---------------------------------------------------------
-# Gradio UI
+# Unified inference router
 # ---------------------------------------------------------
-def inference(model, text, image):
-    if model == "Multilingual Sentiment":
+
+def inference(model_name, text, image):
+    if model_name == "Multilingual Sentiment":
         return run_multilingual(text)
-    if model == "SDG Classification":
+    elif model_name == "SDG Classification":
         return run_sdg(text)
-    if model == "German Sentiment":
+    elif model_name == "German Sentiment":
         return run_german(text)
-    if model == "ViT Image Classification":
+    elif model_name == "ViT Image Classification":
         if image is None:
-            return {"error": "Upload an image"}
+            return {"error": "Please upload an image."}
         return run_vit(image)
-    return {"error": "Invalid selection"}
+    else:
+        return {"error": "Invalid model selected."}
 
 
-with gr.Blocks() as demo:
-    gr.Markdown("# 🔍 Multi-Model ONNX Inference (HF-loaded)")
+# ---------------------------------------------------------
+# Gradio UI
+# ---------------------------------------------------------
 
-    model_choice = gr.Dropdown(
-        ["Multilingual Sentiment", "SDG Classification", "German Sentiment", "ViT Image Classification"],
+with gr.Blocks() as demo:
+    gr.Markdown("# 🔍 Multi-Model ONNX Demo (Loaded from Hugging Face Hub)")
+    gr.Markdown("Text + Image models, running entirely with **ONNX Runtime CPU**")
+
+    model_selector = gr.Dropdown(
+        [
+            "Multilingual Sentiment",
+            "SDG Classification",
+            "German Sentiment",
+            "ViT Image Classification"
+        ],
         label="Choose a model"
     )
 
-    text_in = gr.Textbox(label="Text Input", lines=3)
-    img_in  = gr.Image(label="Image Input", type="pil")
-    output  = gr.JSON(label="Output")
+    text_input = gr.Textbox(lines=3, label="Text Input")
+    image_input = gr.Image(type="pil", label="Image Input")
+    output_box = gr.JSON(label="Model Output")
+
+    run_button = gr.Button("Run")
 
-    run_btn = gr.Button("Run Inference")
-    run_btn.click(fn=inference, inputs=[model_choice, text_in, img_in], outputs=[output])
+    run_button.click(
+        inference,
+        inputs=[model_selector, text_input, image_input],
+        outputs=output_box
+    )
 
 demo.launch()