Update app.py

app.py

@@ -1,21 +1,17 @@
 import gradio as gr
 import nltk
-nltk.download('punkt')
-from lang_list import (
-    LANGUAGE_NAME_TO_CODE,
-    T2TT_TARGET_LANGUAGE_NAMES,
-    TEXT_SOURCE_LANGUAGE_NAMES,
-)
-DEFAULT_TARGET_LANGUAGE = "English"
-from transformers import SeamlessM4TForTextToText
-from transformers import AutoProcessor
-model = SeamlessM4TForTextToText.from_pretrained("facebook/hf-seamless-m4t-medium")
-processor = AutoProcessor.from_pretrained("facebook/hf-seamless-m4t-medium")
+nltk.download('punkt_tab')
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from IndicTransToolkit.IndicTransToolkit import IndicProcessor
+import torch
 
-# text_inputs = processor(text = "Hello, my dog is cute", src_lang="eng", return_tensors="pt")
-# output_tokens = model.generate(**text_inputs, tgt_lang="pan")
-# translated_text_from_text = processor.decode(output_tokens[0].tolist(), skip_special_tokens=True)
-# print(translated_text_from_text)
+# 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)
 
 def split_text_into_batches(text, max_tokens_per_batch):
     sentences = nltk.sent_tokenize(text)  # Tokenize text into sentences
@@ -31,93 +27,88 @@ def split_text_into_batches(text, max_tokens_per_batch):
         batches.append(current_batch.strip())  # Add the last batch
     return batches
 
-def run_t2tt(file_uploader , input_text: str, source_language: str, target_language: str) -> (str, bytes):
+def run_translation(file_uploader, input_text, source_language, target_language):
     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= 256
+        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)
 
-with gr.Blocks() as demo_t2tt:
+    return output, _output_name
+
+# Define Gradio UI
+with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
-            with gr.Group():
-                file_uploader = gr.File(label="Upload a text file (Optional)")
-                input_text = gr.Textbox(label="Input text")
-                with gr.Row():
-                    source_language = gr.Dropdown(
-                        label="Source language",
-                        choices=TEXT_SOURCE_LANGUAGE_NAMES,
-                        value="Punjabi",
-                    )
-                    target_language = gr.Dropdown(
-                        label="Target language",
-                        choices=T2TT_TARGET_LANGUAGE_NAMES,
-                        value=DEFAULT_TARGET_LANGUAGE,
-                    )
+            file_uploader = gr.File(label="Upload a text file (Optional)")
+            input_text = gr.Textbox(label="Input text", lines=5, placeholder="Enter text here...")
+            source_language = gr.Dropdown(
+                label="Source language",
+                choices=["Hindi", "Punjabi", "English"],
+                value="Hindi",
+            )
+            target_language = gr.Dropdown(
+                label="Target language",
+                choices=["Hindi", "Punjabi", "English"],
+                value="English",
+            )
             btn = gr.Button("Translate")
         with gr.Column():
-            output_text = gr.Textbox(label="Translated text")
+            output_text = gr.Textbox(label="Translated text", lines=5)
             output_file = gr.File(label="Translated text file")
 
-    gr.Examples(
-        examples=[
-            [
-                None,
-                "The sinister destruction of the holy Akal Takht and the ruthless massacre of thousands of innocent pilgrims had unmasked the deep-seated hatred and animosity that the Indian Government had been nurturing against Sikhs ever since independence",
-                "English",
-                "Punjabi",
-            ],
-            [
-                None,
-                "It contains. much useful information about administrative, revenue, judicial and ecclesiastical activities in various areas which, it is hoped, would supplement the information available in official records.",
-                "English",
-                "Hindi",
-            ],
-            [
-                None,
-                "दुनिया में बहुत सी अलग-अलग भाषाएं हैं और उनमें अपने वर्ण और शब्दों का भंडार होता है. इसमें में कुछ उनके अपने शब्द होते हैं तो कुछ ऐसे भी हैं, जो दूसरी भाषाओं से लिए जाते हैं.",
-                "Hindi",
-                "Punjabi",
-            ],
-            [
-                None,
-                "ਸੂੂਬੇ ਦੇ ਕਈ ਜ਼ਿਲ੍ਹਿਆਂ ’ਚ ਬੁੱਧਵਾਰ ਸਵੇਰੇ ਸੰਘਣੀ ਧੁੰਦ ਛਾਈ ਰਹੀ ਤੇ ਤੇਜ਼ ਹਵਾਵਾਂ ਨੇ ਕਾਂਬਾ ਹੋਰ ਵਧਾ ਦਿੱਤਾ। ਸੱਤ ਸ਼ਹਿਰਾਂ ’ਚ ਦਿਨ ਦਾ ਤਾਪਮਾਨ ਦਸ ਡਿਗਰੀ ਸੈਲਸੀਅਸ ਦੇ ਆਸਪਾਸ ਰਿਹਾ। ਸੂਬੇ ’ਚ ਵੱਧ ਤੋਂ ਵੱਧ ਤਾਪਮਾਨ ’ਚ ਵੀ ਦਸ ਡਿਗਰੀ ਸੈਲਸੀਅਸ ਦੀ ਗਿਰਾਵਟ ਦਰਜ ਕੀਤੀ ਗਈ",
-                "Punjabi",
-                "English",
-            ],
-        ],
-        inputs=[file_uploader ,input_text, source_language, target_language],
-        outputs=[output_text, output_file],
-        fn=run_t2tt,
-        cache_examples=False,
-        api_name=False,
-    )
-
-    gr.on(
-        triggers=[input_text.submit, btn.click],
-        fn=run_t2tt,
+    btn.click(
+        fn=run_translation,
         inputs=[file_uploader, input_text, source_language, target_language],
         outputs=[output_text, output_file],
-        api_name="t2tt",
     )
 
-with gr.Blocks() as demo:
-    with gr.Tabs():
-        with gr.Tab(label="Translate"):
-            demo_t2tt.render()
-
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(debug=True)