Update app.py

app.py

@@ -4,7 +4,6 @@ from transformers import pipeline
 import os
 import spaces
 
-# Define models for each dialect
 MODELS = {
     "Moroccan": "BounharAbdelaziz/ModernBERT-Arabic-base-stage-3-decay-mx8192-ARY",
     "Arabic": "BounharAbdelaziz/ModernBERT-Arabic-base-stage-3-decay-mx8192-MSA",
@@ -13,7 +12,32 @@ MODELS = {
     "Algerian": "BounharAbdelaziz/ModernBERT-Arabic-base-stage-3-decay-mx8192-ALGERIAN",
 }
 
-TOKEN = os.environ.get("HF_TOKEN", "")
+EXAMPLES = {
+    "Moroccan": [
+        "الدار البيضاء <mask>", 
+        "المغرب بلاد <mask>",
+        "كناكل <mask> فالمغرب",
+        "العاصمة د <mask> هي الرباط",
+        "المغرب <mask> زوين",
+        "انا سميتي مريم، و كنسكن ف<mask> العاصمة دفلسطين"
+    ],
+    "Arabic": [
+        "العاصمة الرسمية لمصر هي <mask>.",
+        "أطول نهر في العالم هو نهر <mask>.",
+        "الشاعر العربي المشهور <mask> كتب قصيدة 'أراك عصي الدمع'.",
+        "عندما أستيقظ في الصباح، أشرب فنجان من <mask>.",
+        "في التأني <mask> وفي العجلة الندامة.",
+        "معركة <mask> كانت من أهم المعارك في تاريخ الإسلام.",
+        "يعتبر <mask> من أهم العلماء في مجال الفيزياء.",
+        "تقع جبال <mask> في شمال إفريقيا.",
+        "يعتبر <mask> من أركان الإسلام الخمسة."
+    ],
+    "Egyptian": ["القاهرة مدينة <mask>", "مصر بلاد <mask>", "بنحب <mask> فمصر"],
+    "Tunisian": ["تونس بلاد <mask>", "المنستير مدينة <mask>", "عيشتي في <mask>"],
+    "Algerian": ["الجزائر بلاد <mask>", "قسنطينة مدينة <mask>", "نحبو <mask> ف الجزائر"],
+}
+
+TOKEN = os.environ["HF_TOKEN"]
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
 def load_model(dialect):
@@ -22,77 +46,94 @@ def load_model(dialect):
 
 pipe = None
 
-@spaces.GPU
+ @spaces.GPU 
 def predict(text, dialect):
     global pipe
-    if pipe is None or dialect != predict.current_dialect:
+    if pipe is None or dialect != predict.current_dialect:  # Reload model if dialect changes
         pipe = load_model(dialect)
         predict.current_dialect = dialect
     outputs = pipe(text)
-    return {x["token_str"]: float(x["score"]) for x in outputs}
+    scores = [x["score"] for x in outputs]
+    tokens = [x["token_str"] for x in outputs]
+    return {label: float(prob) for label, prob in zip(tokens, scores)}
 
+# Initialize current dialect
 predict.current_dialect = None
 
-# Create the app
+# Create Gradio interface
 with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
+            # Dropdown for dialect selection
             dialect_dropdown = gr.Dropdown(
                 choices=["Arabic", "Tunisian", "Moroccan", "Algerian", "Egyptian"],
-                value="Arabic",
-                label="Select Dialect"
+                label="Select Dialect",
+                value="Arabic"
             )
             
+            # Input text box
             input_text = gr.Textbox(
                 label="Input",
-                placeholder="Enter text with <mask>...",
+                placeholder="Enter text here...",
                 rtl=True
             )
             
-            # Use a HTML component to display clickable examples
-            examples_html = gr.HTML(
-                value="<div id='examples'>Examples: " + 
-                      " | ".join([f"<a href='#' onclick='return false;'>{ex}</a>" for ex in EXAMPLES["Arabic"]]) + 
-                      "</div>"
-            )
-            
+            # Button row
             with gr.Row():
                 clear_btn = gr.Button("Clear")
                 submit_btn = gr.Button("Submit", variant="primary")
+            
+            # Examples section (stateful)
+            examples_state = gr.State(EXAMPLES["Arabic"])  # Store examples for the current dialect
+            
+            example_dropdown = gr.Dropdown(
+                choices=EXAMPLES["Arabic"],
+                label="Select Example",
+                interactive=True
+            )
+            
+            load_example_btn = gr.Button("Load Example")
 
         with gr.Column():
-            output_labels = gr.Label(label="Predictions")
+            # Output probabilities
+            output_labels = gr.Label(
+                label="Prediction Results",
+                show_label=False
+            )
 
-    # JavaScript to handle example clicks
-    demo.load(
-        None,
-        None,
-        None,
-        _js="""
-        () => {
-            document.querySelectorAll('#examples a').forEach(a => {
-                a.addEventListener('click', (e) => {
-                    const text = e.target.textContent;
-                    document.querySelector('textarea[data-testid="textbox"]').value = text;
-                });
-            });
-        }
-        """
-    )
+    # Function to update examples when dialect changes
+    def update_example_choices(dialect):
+        return gr.update(choices=EXAMPLES.get(dialect, EXAMPLES["Arabic"]), value=EXAMPLES[dialect][0])
 
-    # Update examples when dialect changes
-    def update_examples(dialect):
-        examples = EXAMPLES.get(dialect, EXAMPLES["Arabic"])
-        html = "<div id='examples'>Examples: " + " | ".join([f"<a href='#' onclick='return false;'>{ex}</a>" for ex in examples]) + "</div>"
-        return gr.HTML.update(value=html)
+    # Function to load selected example into the text box
+    def load_example(selected_example):
+        return selected_example
 
+    # Update example dropdown choices dynamically
     dialect_dropdown.change(
-        update_examples,
+        update_example_choices,
         inputs=dialect_dropdown,
-        outputs=examples_html
+        outputs=example_dropdown
     )
 
-    submit_btn.click(predict, inputs=[input_text, dialect_dropdown], outputs=output_labels)
-    clear_btn.click(lambda: "", outputs=input_text)
+    # Load selected example into the input text box
+    load_example_btn.click(
+        load_example,
+        inputs=example_dropdown,
+        outputs=input_text
+    )
+
+    # Button actions
+    submit_btn.click(
+        predict,
+        inputs=[input_text, dialect_dropdown],
+        outputs=output_labels
+    )
+    
+    clear_btn.click(
+        lambda: "",
+        outputs=input_text
+    )
 
+# Launch the app
 demo.launch()