Update app.py

app.py

@@ -6,26 +6,37 @@ from transformers import MarianMTModel, MarianTokenizer
 import os
 os.system("pip install sentencepiece")
 
-# Load models and tokenizers
-model_en_to_ur = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-en-ur")
+# Check if GPU is available and use it
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load models and tokenizers once (globally)
+model_en_to_ur = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-en-ur").to(device)
 tokenizer_en_to_ur = MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-ur")
 
-model_ur_to_en = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-ur-en")
+model_ur_to_en = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-ur-en").to(device)
 tokenizer_ur_to_en = MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-ur-en")
 
+# Apply torch.compile() for optimization (if using PyTorch 2.0+)
+if torch.__version__ >= "2.0":
+    model_en_to_ur = torch.compile(model_en_to_ur)
+    model_ur_to_en = torch.compile(model_ur_to_en)
+
 # Function to translate text
 def translate(text, direction):
+    if not text.strip():
+        return "Please enter some text to translate."
+
     if direction == "English to Urdu":
         tokenizer, model = tokenizer_en_to_ur, model_en_to_ur
     else:
         tokenizer, model = tokenizer_ur_to_en, model_ur_to_en
 
-    # Tokenize input text
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+    # Tokenize input text (optimized padding)
+    inputs = tokenizer(text, return_tensors="pt", padding="longest", truncation=True).to(device)
 
     # Generate translation
     with torch.no_grad():
-        translated = model.generate(**inputs)
+        translated = model.generate(**inputs, max_length=512)
 
     # Decode output text
     output = tokenizer.decode(translated[0], skip_special_tokens=True)
@@ -39,8 +50,8 @@ interface = gr.Interface(
         gr.Radio(["English to Urdu", "Urdu to English"], label="Translation Direction", value="English to Urdu"),
     ],
     outputs=gr.Textbox(label="Translated Text"),
-    title="English ↔ Urdu Translator",
-    description="Translate text between English and Urdu using a neural machine translation model.",
+    title="⚡ Fast English ↔ Urdu Translator",
+    description="Translate text between English and Urdu quickly using a neural machine translation model with GPU acceleration.",
 )
 
 # Launch the Gradio app