Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
-import tiktoken
+import torch
+
 # Load the model and tokenizer
 model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
 tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
@@ -35,35 +36,51 @@ language_codes = {
 }
 
 def translate(text, src_lang, tgt_lang):
-    # Set the source language
-    tokenizer.src_lang = language_codes[src_lang]
-    
-    # Tokenize the input text
-    encoded = tokenizer(text, return_tensors="pt")
-    
-    # Generate translation
-    generated_tokens = model.generate(
-        **encoded,
-        forced_bos_token_id=tokenizer.lang_code_to_id[language_codes[tgt_lang]]
-    )
-    
-    # Decode the generated tokens
-    translation = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
-    
-    return translation
+    try:
+        if not text.strip():
+            return "Please enter some text to translate."
+        
+        if src_lang == tgt_lang:
+            return text
+        
+        # Set the source language
+        tokenizer.src_lang = language_codes[src_lang]
+        
+        # Tokenize the input text
+        encoded = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
+        
+        # Generate translation
+        with torch.no_grad():
+            generated_tokens = model.generate(
+                **encoded,
+                forced_bos_token_id=tokenizer.lang_code_to_id[language_codes[tgt_lang]],
+                max_length=512,
+                num_beams=5,
+                length_penalty=1.0
+            )
+        
+        # Decode the generated tokens
+        translation = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+        
+        return translation
+    except Exception as e:
+        return f"Translation error: {str(e)}"
 
 # Create the Gradio interface
-iface = gr.Interface(
+demo = gr.Interface(
     fn=translate,
     inputs=[
-        gr.Textbox(label="Input Text"),
-        gr.Dropdown(choices=list(language_codes.keys()), label="Source Language"),
-        gr.Dropdown(choices=list(language_codes.keys()), label="Target Language"),
+        gr.Textbox(label="Input Text", placeholder="Enter text to translate..."),
+        gr.Dropdown(choices=sorted(language_codes.keys()), label="Source Language", value="English"),
+        gr.Dropdown(choices=sorted(language_codes.keys()), label="Target Language", value="Spanish"),
     ],
     outputs=gr.Textbox(label="Translated Text"),
     title="Multilingual Translation with MBart",
     description="Translate text between multiple languages using the MBart model.",
+    examples=[
+        ["Hello, how are you?", "English", "Spanish"],
+        ["Bonjour, comment allez-vous?", "French", "English"],
+    ]
 )
 
-# Launch the interface
-iface.launch()
+demo.launch()