app.py

app.py

@@ -0,0 +1,104 @@
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
+
+# Load the model and tokenizer
+model_name = "jbochi/madlad400-3b-mt"
+
+@gr.cache_resource()
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSeq2SeqLM.from_pretrained(
+        model_name,
+        torch_dtype=torch.float16,
+        device_map="auto"
+    )
+    return tokenizer, model
+
+tokenizer, model = load_model()
+
+def translate_text(text, source_lang, target_lang):
+    """
+    Translate text between English and Persian using MADLAD-400-3B
+    """
+    lang_codes = {"English": "en", "Persian": "fa"}
+    source_code = lang_codes[source_lang]
+    target_code = lang_codes[target_lang]
+    
+    prompt = f"<2{target_code}> {text}"
+    
+    try:
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
+        inputs = {k: v.to(model.device) for k, v in inputs.items()}
+        
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_length=512,
+                num_beams=5,
+                early_stopping=True,
+                no_repeat_ngram_size=3,
+            )
+        
+        translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return translated_text
+    
+    except Exception as e:
+        return f"Error during translation: {str(e)}"
+
+with gr.Blocks(title="English-Persian Translator", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🌍 English-Persian Translator")
+    gr.Markdown("Translate text between English and Persian using MADLAD-400-3B model")
+    
+    with gr.Row():
+        with gr.Column():
+            source_lang = gr.Dropdown(
+                choices=["English", "Persian"],
+                value="English",
+                label="From"
+            )
+            input_text = gr.Textbox(
+                lines=4,
+                placeholder="Enter text to translate...",
+                label="Input Text"
+            )
+        
+        with gr.Column():
+            target_lang = gr.Dropdown(
+                choices=["Persian", "English"],
+                value="Persian",
+                label="To"
+            )
+            output_text = gr.Textbox(
+                lines=4,
+                label="Translated Text",
+                interactive=False
+            )
+    
+    translate_btn = gr.Button("Translate ✨", variant="primary")
+    
+    gr.Examples(
+        examples=[
+            ["Hello, how are you today?", "English", "Persian"],
+            ["What is your name?", "English", "Persian"],
+            ["سلام، حالتون چطوره؟", "Persian", "English"],
+            ["امروز هوا خوب است", "Persian", "English"]
+        ],
+        inputs=[input_text, source_lang, target_lang],
+        outputs=output_text,
+        fn=translate_text
+    )
+    
+    translate_btn.click(
+        fn=translate_text,
+        inputs=[input_text, source_lang, target_lang],
+        outputs=output_text
+    )
+    
+    def update_target_lang(source_lang):
+        return "Persian" if source_lang == "English" else "English"
+    
+    source_lang.change(update_target_lang, source_lang, target_lang)
+
+if __name__ == "__main__":
+    demo.launch()