Remove max tokens slider and simplify interface

app.py

@@ -1,55 +1,49 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
 # Load both translation models from Hugging Face
 # English to Moroccan Arabic (Darija)
-tokenizer_eng_to_darija = AutoTokenizer.from_pretrained("Saidtaoussi/AraT5_Darija_to_MSA")
-model_eng_to_darija = AutoModelForSeq2SeqLM.from_pretrained("Saidtaoussi/AraT5_Darija_to_MSA")
+tokenizer_eng_to_darija = AutoTokenizer.from_pretrained("lachkarsalim/Helsinki-translation-English_Moroccan-Arabic")
+model_eng_to_darija = AutoModelForSeq2SeqLM.from_pretrained("lachkarsalim/Helsinki-translation-English_Moroccan-Arabic")
 
 # Moroccan Arabic (Darija) to Modern Standard Arabic (MSA)
-tokenizer_darija_to_msa = AutoTokenizer.from_pretrained("lachkarsalim/Helsinki-translation-English_Moroccan-Arabic")
-model_darija_to_msa = AutoModelForSeq2SeqLM.from_pretrained("lachkarsalim/Helsinki-translation-English_Moroccan-Arabic")
-
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    translation_choice: str,
-):
-    # Ensure there's no empty input
-    if not message.strip():
-        return "Error: Please enter a valid text to translate."
-
-    # Initialize the response variable
-    response = ""
-
+tokenizer_darija_to_msa = AutoTokenizer.from_pretrained("Saidtaoussi/AraT5_Darija_to_MSA")
+model_darija_to_msa = AutoModelForSeq2SeqLM.from_pretrained("Saidtaoussi/AraT5_Darija_to_MSA")
+
+# Translation function for Darija to MSA
+def translate_darija_to_msa(darija_text):
+    inputs = tokenizer_darija_to_msa(darija_text, return_tensors="pt", padding=True)
+    translated = model_darija_to_msa.generate(**inputs)
+    translated_text = tokenizer_darija_to_msa.decode(translated[0], skip_special_tokens=True)
+    return translated_text
+
+# Translation function for English to Moroccan Arabic and vice versa
+def translate_eng_to_darija(eng_text, direction="eng_to_darija"):
+    if direction == "eng_to_darija":
+        inputs = tokenizer_eng_to_darija(eng_text, return_tensors="pt", padding=True)
+        translated = model_eng_to_darija.generate(**inputs)
+        translated_text = tokenizer_eng_to_darija.decode(translated[0], skip_special_tokens=True)
+    else:
+        # Translate from Darija to English (reverse translation)
+        inputs = tokenizer_eng_to_darija(eng_text, return_tensors="pt", padding=True)
+        translated = model_eng_to_darija.generate(**inputs)
+        translated_text = tokenizer_eng_to_darija.decode(translated[0], skip_special_tokens=True)
+    return translated_text
+
+
+# Gradio interface setup without max new tokens
+def respond(message, translation_choice: str):
     # Translate based on the user's choice
-    try:
-        if translation_choice == "Moroccan Arabic to MSA":
-            # Translate Moroccan Arabic (Darija) to Modern Standard Arabic
-            inputs = tokenizer_darija_to_msa(message, return_tensors="pt", padding=True)
-            outputs = model_darija_to_msa.generate(inputs["input_ids"], num_beams=5, max_length=max_tokens, early_stopping=True)
-            response = tokenizer_darija_to_msa.decode(outputs[0], skip_special_tokens=True)
-
-        elif translation_choice == "English to Moroccan Arabic":
-            # Translate English to Moroccan Arabic (Darija)
-            inputs = tokenizer_eng_to_darija(message, return_tensors="pt", padding=True)
-            outputs = model_eng_to_darija.generate(inputs["input_ids"], num_beams=5, max_length=max_tokens, early_stopping=True)
-            response = tokenizer_eng_to_darija.decode(outputs[0], skip_special_tokens=True)
-    except Exception as e:
-        response = f"Error occurred: {str(e)}"
+    if translation_choice == "Moroccan Arabic to MSA":
+        return translate_darija_to_msa(message)
+    elif translation_choice == "English to Moroccan Arabic":
+        return translate_eng_to_darija(message, direction="eng_to_darija")
 
-    return response
 
-
-# Gradio interface setup without pre-filled system message
 demo = gr.Interface(
     fn=respond,
     inputs=[
         gr.Textbox(value="", label="Enter Your Text", placeholder="Type your sentence here..."),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
         gr.Dropdown(
             label="Choose Translation Direction",
             choices=["English to Moroccan Arabic", "Moroccan Arabic to MSA"],
@@ -61,3 +55,4 @@ demo = gr.Interface(
 
 # Launch the interface
 demo.launch()
+