Update translation.py

translation.py

@@ -56,7 +56,7 @@ class CombinedModel:
         input_ids = kwargs.get('input_ids')
         if not input_ids:
             return torch.tensor([])
-        return torch.tensor([combined_translate(tokenizer.decode(x, skip_special_tokens=True), self.source_lang, self.target_lang, self.default_tokenizer, self.default_model) for x in input_ids])
+        return torch.tensor([combined_translate(self.default_tokenizer.decode(x, skip_special_tokens=True), self.source_lang, self.target_lang, self.default_tokenizer, self.default_model) for x in input_ids])
 
 # Function to load appropriate translation model with optimized caching
 @st.cache_resource
@@ -93,8 +93,10 @@ def translate(text, source_lang, target_lang):
     try:
         tokenizer, model = load_model(source_lang, target_lang)
         inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=500)
+        if inputs['input_ids'].size(0) > 1:  # Ensure single sequence
+            inputs = {k: v[0].unsqueeze(0) for k, v in inputs.items()}
         with torch.no_grad():
-            translated = model.generate(**inputs, max_length=1000 if target_lang == "hi" else 500, num_beams=4, early_stopping=True)  # Reduced to 4 beams for speed
+            translated = model.generate(**inputs, max_length=1000 if target_lang == "ja" else 500, num_beams=4, early_stopping=True)
         result = tokenizer.decode(translated[0], skip_special_tokens=True)
         return result if result.strip() else text
     except Exception as e: