Update app.py

app.py

@@ -14,19 +14,21 @@ from happytransformer import HappyTextToText, TTSettings
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM,logging
 from transformers.integrations import deepspeed
 import re
+from IndicTransToolkit import IndicProcessor
 import torch
-from lang_list import (
-    LANGUAGE_NAME_TO_CODE,
-    T2TT_TARGET_LANGUAGE_NAMES,
-    TEXT_SOURCE_LANGUAGE_NAMES,
-)
+import torch
+
 logging.set_verbosity_error()
 
 DEFAULT_TARGET_LANGUAGE = "English"
-from transformers import SeamlessM4TForTextToText
-from transformers import AutoProcessor
-model = SeamlessM4TForTextToText.from_pretrained("facebook/hf-seamless-m4t-large")
-processor = AutoProcessor.from_pretrained("facebook/hf-seamless-m4t-large")
+# Load IndicTrans2 model
+model_name = "ai4bharat/indictrans2-indic-indic-dist-320M"
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name, trust_remote_code=True)
+ip = IndicProcessor(inference=True)
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(DEVICE)
+
 
 
 import pytesseract as pt
@@ -174,22 +176,59 @@ def split_text_into_batches(text, max_tokens_per_batch):
 @spaces.GPU(duration=60)
 def run_t2tt(file_uploader , input_text: str, source_language: str, target_language: str) -> (str, bytes):
     if file_uploader is not None:
-        with open(file_uploader, 'r') as file:
-            input_text=file.read()
-    source_language_code = LANGUAGE_NAME_TO_CODE[source_language]
-    target_language_code = LANGUAGE_NAME_TO_CODE[target_language]
-    max_tokens_per_batch= 2048
+        with open(file_uploader.name, "r", encoding="utf-8") as file:
+            input_text = file.read()
+
+    # Language mapping
+    lang_code_map = {
+        "Hindi": "hin_Deva",
+        "Punjabi": "pan_Guru",
+        "English": "eng_Latn",
+    }
+
+    src_lang = lang_code_map[source_language]
+    tgt_lang = lang_code_map[target_language]
+
+    max_tokens_per_batch = 256
     batches = split_text_into_batches(input_text, max_tokens_per_batch)
     translated_text = ""
+
     for batch in batches:
-        text_inputs = processor(text=batch, src_lang=source_language_code, return_tensors="pt")
-        output_tokens = model.generate(**text_inputs, tgt_lang=target_language_code)
-        translated_batch = processor.decode(output_tokens[0].tolist(), skip_special_tokens=True)
-        translated_text += translated_batch + " "
-    output=translated_text.strip()
+        batch_preprocessed = ip.preprocess_batch([batch], src_lang=src_lang, tgt_lang=tgt_lang)
+        inputs = tokenizer(
+            batch_preprocessed,
+            truncation=True,
+            padding="longest",
+            return_tensors="pt",
+            return_attention_mask=True,
+        ).to(DEVICE)
+
+        with torch.no_grad():
+            generated_tokens = model.generate(
+                **inputs,
+                use_cache=True,
+                min_length=0,
+                max_length=256,
+                num_beams=5,
+                num_return_sequences=1,
+            )
+
+        with tokenizer.as_target_tokenizer():
+            decoded_tokens = tokenizer.batch_decode(
+                generated_tokens.detach().cpu().tolist(),
+                skip_special_tokens=True,
+                clean_up_tokenization_spaces=True,
+            )
+
+        translations = ip.postprocess_batch(decoded_tokens, lang=tgt_lang)
+        translated_text += " ".join(translations) + " "
+
+    output = translated_text.strip()
     _output_name = "result.txt"
-    open(_output_name, 'w').write(output)
-    return str(output), _output_name
+    with open(_output_name, "w", encoding="utf-8") as out_file:
+        out_file.write(output)
+
+    return output, _output_name
 
 with gr.Blocks() as demo_t2tt:
     with gr.Row():