Update app.py

app.py

@@ -9,7 +9,7 @@ MODEL_NAME = "Abelex/afro-xlmr-large"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 # --------------------------------------------------
-# Load tokenizer & model (CRITICAL FIX)
+# Load tokenizer & model
 # --------------------------------------------------
 tokenizer = AutoTokenizer.from_pretrained(
     MODEL_NAME,
@@ -28,8 +28,8 @@ model.eval()
 # Prediction function
 # --------------------------------------------------
 def classify_text(text):
-    if not text.strip():
-        return "⚠️ Please enter Amharic text", {}
+    if not text or not text.strip():
+        return gr.Markdown("⚠️ **Please enter Amharic text**"), None
 
     inputs = tokenizer(
         text,
@@ -41,50 +41,54 @@ def classify_text(text):
 
     with torch.no_grad():
         outputs = model(**inputs)
-
-        # 🔑 Handle custom forward outputs safely
-        if hasattr(outputs, "logits"):
-            logits = outputs.logits
-        else:
-            logits = outputs[0]
-
+        logits = outputs.logits if hasattr(outputs, "logits") else outputs[0]
         probs = torch.softmax(logits, dim=-1)[0]
 
     pred_id = torch.argmax(probs).item()
-
-    # Safe label handling
-    id2label = getattr(model.config, "id2label", None)
-    if id2label:
-        pred_label = id2label.get(pred_id, str(pred_id))
-    else:
-        pred_label = f"Class {pred_id}"
+    id2label = getattr(model.config, "id2label", {})
+    pred_label = id2label.get(pred_id, f"Class {pred_id}")
 
     scores = {
-        (id2label.get(i, f"Class {i}") if id2label else f"Class {i}"): float(probs[i])
+        id2label.get(i, f"Class {i}"): float(probs[i])
         for i in range(len(probs))
     }
 
-    return pred_label, scores
+    return (
+        gr.Markdown(f"### 🏷️ **{pred_label}**"),
+        scores
+    )
 
 # --------------------------------------------------
-# Gradio UI
+# Gradio UI (MINIMAL & ATTRACTIVE)
 # --------------------------------------------------
-with gr.Blocks(title="Amharic Text Classification") as demo:
-    gr.Markdown("""
-    ## 📄 Amharic Text Classification  
-    **Model:** Abelex/Sentence-Chunking-Afri_BERTA_amharic_longtext  
-    """)
-
-    input_text = gr.Textbox(
-        lines=8,
-        label="Input Amharic Text",
-        placeholder="እባክዎ የአማርኛ ጽሑፍ እዚህ ያስገቡ..."
+with gr.Blocks(
+    title="Amharic Text Classification",
+    theme=gr.themes.Soft()
+) as demo:
+
+    gr.Markdown(
+        """
+        # 🇪🇹 Amharic Text Classification  
+        <small>Powered by **Afro-XLMR / AfriBERTa**</small>
+        """,
+        elem_id="title"
     )
 
-    classify_btn = gr.Button("🔍 Classify")
+    with gr.Column(scale=1):
+        input_text = gr.Textbox(
+            lines=5,
+            placeholder="እባክዎ የአማርኛ ጽሑፍ እዚህ ያስገቡ...",
+            show_label=False
+        )
+
+        classify_btn = gr.Button(
+            "Classify",
+            variant="primary"
+        )
 
-    output_label = gr.Label(label="Predicted Label")
-    output_scores = gr.JSON(label="Class Probabilities")
+    with gr.Column():
+        output_label = gr.Markdown()
+        output_scores = gr.JSON(label="Class Probabilities")
 
     classify_btn.click(
         fn=classify_text,
@@ -92,7 +96,24 @@ with gr.Blocks(title="Amharic Text Classification") as demo:
         outputs=[output_label, output_scores]
     )
 
-    gr.Markdown("---\nBuilt with ❤️ using Gradio & Hugging Face")
+    gr.Examples(
+        examples=[
+            ["የኢትዮጵያ ኢኮኖሚ በአዲስ እቅድ እየተሻሻለ ነው።"],
+            ["የእግር ኳስ ውድድር በአዲስ መልክ ተጀመረ።"]
+        ],
+        inputs=input_text,
+        label="Examples"
+    )
+
+    gr.Markdown(
+        """
+        ---
+        <small>
+        Model: <b>Abelex/afro-xlmr-large</b><br>
+        Built with ❤️ using Gradio
+        </small>
+        """
+    )
 
 # --------------------------------------------------
 # Launch