Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoConfig
 import torch
 
 # Load model and tokenizer from Hugging Face Hub
@@ -7,8 +7,15 @@ model_name = "humy65/hebrew-intent-classifier"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
 
-# Define intent labels (adjust to match your model's output)
-intent_labels = ["שכחת סיסמה", "ביטול מנוי", "שאלה כללית", "תמיכה טכנית"]
+# Load config and extract label mapping
+config = AutoConfig.from_pretrained(model_name)
+
+# Try to get label names from config.id2label (if available)
+if hasattr(config, "id2label") and config.id2label:
+    intent_labels = [config.id2label[i] for i in range(config.num_labels)]
+else:
+    # Fallback: generate generic labels
+    intent_labels = [f"כוונה {i}" for i in range(config.num_labels)]
 
 def classify_intent(text):
     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
@@ -16,18 +23,25 @@ def classify_intent(text):
         outputs = model(**inputs)
         probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
         top_prob, top_label = torch.max(probs, dim=1)
-        intent = intent_labels[top_label.item()]
+        try:
+            intent = intent_labels[top_label.item()]
+        except IndexError:
+            intent = "כוונה לא מזוהה"
         confidence = top_prob.item()
     return f"כוונה: {intent}\nרמת ביטחון: {confidence:.2f}"
 
 # Build Gradio interface
-demo = gr.Interface(
-    fn=classify_intent,
-    inputs=gr.Textbox(label="הקלד שאלה בעברית", placeholder="לדוגמה: שכחתי את הסיסמה שלי"),
-    outputs=gr.Textbox(label="תוצאה"),
-    title="🔍 מסווג כוונות בעברית",
-    description="מודל לזיהוי כוונות בשפה העברית. הקלד שאלה וראה את הכוונה המשוערת.",
-    theme="soft"
-)
+with gr.Blocks(theme="soft") as demo:
+    gr.Markdown("## 🔍 מסווג כוונות בעברית")
+    gr.Markdown("מודל לזיהוי כוונות בשפה העברית. הקלד שאלה וראה את הכוונה המשוערת.")
+    
+    with gr.Row():
+        input_text = gr.Textbox(label="הקלד שאלה בעברית", placeholder="לדוגמה: שכחתי את הסיסמה שלי")
+        intent_list = gr.Dropdown(label="כוונות זמינות", choices=intent_labels, interactive=False)
+    
+    output_text = gr.Textbox(label="תוצאה")
+    submit_btn = gr.Button("סווג")
+
+    submit_btn.click(fn=classify_intent, inputs=input_text, outputs=output_text)
 
 demo.launch()