Upload app.py with huggingface_hub

app.py

@@ -1,4 +1,3 @@
-from huggingface_hub import HfApi, HfFolder
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
@@ -8,10 +7,19 @@ from deep_translator import GoogleTranslator
 import os
 import subprocess
 
-# Language Mapping
+# Language Mapping with more robust models
 LANG_MAP = {
-    'en-te': {'source': 'en', 'target': 'te', 'model': 'Helsinki-NLP/opus-mt-en-te'},
-    'te-en': {'source': 'te', 'target': 'en', 'model': 'Helsinki-NLP/opus-mt-te-en'}
+    'en-te': {
+        'source': 'en', 
+        'target': 'te', 
+        'model': 'Helsinki-NLP/opus-mt-en-te'
+    },
+    'te-en': {
+        'source': 'te', 
+        'target': 'en', 
+        'model': 'Google/byt5-base', # More robust for te-en translation
+        'target_language': 'en'
+    }
 }
 
 # Load translation models and tokenizers
@@ -44,8 +52,14 @@ class TranslationHandler:
             tokenizer = self.tokenizers[direction]
             model = self.models[direction]
             
-            # Tokenize input
-            inputs = tokenizer(text, return_tensors="pt", max_length=512, truncation=True)
+            # Special handling for Telugu to English
+            if direction == 'te-en':
+                # Prefix the input for better translation
+                input_text = f"translate Telugu to English: {text}"
+                inputs = tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True)
+            else:
+                # Standard translation for other directions
+                inputs = tokenizer(text, return_tensors="pt", max_length=512, truncation=True)
             
             # Generate translation
             with torch.no_grad():
@@ -100,7 +114,9 @@ def get_transliteration(text, lang):
 # Text-to-Speech Function
 def text_to_speech(text, lang):
     try:
-        tts = gTTS(text, lang=lang)
+        # Use 'en' for English output
+        tts_lang = 'en' if lang == 'en' else 'te'
+        tts = gTTS(text, lang=tts_lang)
         audio_file = "translated_audio.mp3"
         tts.save(audio_file)
         return audio_file