Added all supported languages

app.py

@@ -4,36 +4,67 @@ from transformers import pipeline
 import os
 import spaces
 
-MODEL_PATH = "BounharAbdelaziz/ModernBERT-Arabic-base-stage-3-decay-mx8192-ARY"
-TOKEN=os.environ["HF_TOKEN"]
-device = "cuda:0" if torch.cuda.is_available() else "cpu"
+# Define models for each dialect
+MODELS = {
+    "Arabic": "BounharAbdelaziz/ModernBERT-Arabic-base-stage-3-decay-mx8192-MSA",
+    "Tunisian": "BounharAbdelaziz/ModernBERT-Arabic-base-stage-3-decay-mx8192-ARY"
+    "Moroccan": "BounharAbdelaziz/ModernBERT-Arabic-base-stage-3-decay-mx8192-ARY"
+    "Algerian": "BounharAbdelaziz/ModernBERT-Arabic-base-stage-3-decay-mx8192-ALGERIAN"
+    "Egyptian": "BounharAbdelaziz/ModernBERT-Arabic-base-stage-3-decay-mx8192-ARZ"
+}
 
+TOKEN = os.environ["HF_TOKEN"]
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
-def load_model():
-    pipe=pipeline(
+# Load the pipeline dynamically based on the selected dialect
+def load_model(dialect):
+    model_path = MODELS.get(dialect, MODELS["Arabic"])  # Default to Arabic if dialect not found
+    pipe = pipeline(
         task="fill-mask",
-        model=MODEL_PATH,
+        model=model_path,
         token=TOKEN,
         device=device
     )
     return pipe
 
-print("[INFO] Lading pipeline model ...")
-pipe=load_model()
-print("[INFO] Pipeline loaded!")
+# Store the pipeline globally
+pipe = None
 
+# Define examples for each dialect
+EXAMPLES = {
+    "Moroccan": ["الدار البيضاء [MASK]", "المغرب بلاد [MASK]", "كناكل [MASK] فالمغرب", "العاصمة د [MASK] هي الرباط","المغرب [MASK] زوين","انا سميتي مريم، و كنسكن ف[MASK] العاصمة دفلسطين"],
+    "Arabic": ["العاصمة د [MASK] هي الرباط", "المغرب [MASK] زوين", "انا سميتي مريم، و كنسكن ف[MASK] العاصمة دفلسطين"],
+    "Egyptian": ["القاهرة مدينة [MASK]", "مصر بلاد [MASK]", "بنحب [MASK] فمصر"],
+    "Tunisian": ["تونس بلاد [MASK]", "المنستير مدينة [MASK]", "عيشتي في [MASK]"],
+    "Algerian": ["الجزائر بلاد [MASK]", "قسنطينة مدينة [MASK]", "نحبو [MASK] ف الجزائر"],
+}
 
+# Predict function
 @spaces.GPU
-def predict(text):
+def predict(text, dialect):
+    global pipe
+    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)
-    scores= [x["score"] for x in outputs]
-    tokens= [x["token_str"] 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 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"],
+                label="Select Dialect",
+                value="Arabic"
+            )
+            
             # Input text box
             input_text = gr.Textbox(
                 label="Input",
@@ -46,10 +77,11 @@ with gr.Blocks() as demo:
                 clear_btn = gr.Button("Clear")
                 submit_btn = gr.Button("Submit", variant="primary")
             
-            # Examples section
-            gr.Examples(
-                examples=["العاصمة د <mask> هي الرباط","المغرب <mask> زوين","انا سميتي مريم، و كنسكن ف<mask> العاصمة دفلسطين"],
-                inputs=input_text
+            # Examples section (dynamic based on dialect)
+            examples = gr.Examples(
+                examples=EXAMPLES["Arabic"],  # Default to Arabic examples
+                inputs=input_text,
+                label="Examples"
             )
         
         with gr.Column():
@@ -59,10 +91,20 @@ with gr.Blocks() as demo:
                 show_label=False
             )
 
+    # Update examples when dialect changes
+    def update_examples(dialect):
+        return EXAMPLES.get(dialect, EXAMPLES["Arabic"])
+
+    dialect_dropdown.change(
+        update_examples,
+        inputs=dialect_dropdown,
+        outputs=examples
+    )
+
     # Button actions
     submit_btn.click(
         predict,
-        inputs=input_text,
+        inputs=[input_text, dialect_dropdown],
         outputs=output_labels
     )