Update app.py

app.py

@@ -1,49 +1,44 @@
+import gradio as gr
 import torch
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from transformers import (
+    AutoModelForSeq2SeqLM,
+    AutoTokenizer,
+)
 from IndicTransToolkit import IndicProcessor
-import gradio as gr
+import os
+import subprocess
 
-# Define the model and tokenizer
-model_name = "ai4bharat/indictrans2-indic-indic-1B"
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name, trust_remote_code=True)
-ip = IndicProcessor(inference=True)
+# Function to clone the repository and set up the environment
+def setup_repo():
+    # Clone the repository
+    repo_url = "https://github.com/AI4Bharat/IndicTrans2"
+    repo_dir = "IndicTrans2"
+    
+    if not os.path.exists(repo_dir):
+        subprocess.run(["git", "clone", repo_url])
+    
+    # Navigate to the project directory and install dependencies
+    os.chdir(os.path.join(repo_dir, "huggingface_interface"))
+    subprocess.run(["source", "install.sh"], shell=True)
 
-# Define the language codes
-LANGUAGES = {
-    "Assamese (asm_Beng)": "asm_Beng",
-    "Kashmiri (kas_Arab)": "kas_Arab",
-    "Punjabi (pan_Guru)": "pan_Guru",
-    "Bengali (ben_Beng)": "ben_Beng",
-    "Kashmiri (kas_Deva)": "kas_Deva",
-    "Sanskrit (san_Deva)": "san_Deva",
-    "Bodo (brx_Deva)": "brx_Deva",
-    "Maithili (mai_Deva)": "mai_Deva",
-    "Santali (sat_Olck)": "sat_Olck",
-    "Dogri (doi_Deva)": "doi_Deva",
-    "Malayalam (mal_Mlym)": "mal_Mlym",
-    "Sindhi (snd_Arab)": "snd_Arab",
-    "English (eng_Latn)": "eng_Latn",
-    "Marathi (mar_Deva)": "mar_Deva",
-    "Sindhi (snd_Deva)": "snd_Deva",
-    "Konkani (gom_Deva)": "gom_Deva",
-    "Manipuri (mni_Beng)": "mni_Beng",
-    "Tamil (tam_Taml)": "tam_Taml",
-    "Gujarati (guj_Gujr)": "guj_Gujr",
-    "Manipuri (mni_Mtei)": "mni_Mtei",
-    "Telugu (tel_Telu)": "tel_Telu",
-    "Hindi (hin_Deva)": "hin_Deva",
-    "Nepali (npi_Deva)": "npi_Deva",
-    "Urdu (urd_Arab)": "urd_Arab",
-    "Kannada (kan_Knda)": "kan_Knda",
-    "Odia (ory_Orya)": "ory_Orya",
-}
-
-# Define the translation function
-def translate(text, src_lang, tgt_lang):
-    batch = ip.preprocess_batch([text], src_lang=src_lang, tgt_lang=tgt_lang)
+# Function to process translation
+def translate(input_text, src_lang, tgt_lang):
+    setup_repo()  # Ensure the repo is set up
+    model_name = "ai4bharat/indictrans2-indic-indic-1B"
+    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name, trust_remote_code=True)
+    ip = IndicProcessor(inference=True)
+    
+    batch = ip.preprocess_batch([input_text], src_lang=src_lang, tgt_lang=tgt_lang)
     DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-    inputs = tokenizer(batch, truncation=True, padding="longest", return_tensors="pt").to(DEVICE)
+    inputs = tokenizer(
+        batch,
+        truncation=True,
+        padding="longest",
+        return_tensors="pt",
+        return_attention_mask=True,
+    ).to(DEVICE)
+
     with torch.no_grad():
         generated_tokens = model.generate(
             **inputs,
@@ -53,22 +48,42 @@ def translate(text, src_lang, tgt_lang):
             num_beams=5,
             num_return_sequences=1,
         )
+
     with tokenizer.as_target_tokenizer():
-        generated_text = tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
-    return generated_text
+        translation = tokenizer.batch_decode(
+            generated_tokens.detach().cpu().tolist(),
+            skip_special_tokens=True,
+            clean_up_tokenization_spaces=True,
+        )[0]
+
+    return translation
+
+# List of languages with their code names
+languages = [
+    ("Assamese", "asm_Beng"), ("Kashmiri (Arabic)", "kas_Arab"), ("Punjabi", "pan_Guru"),
+    ("Bengali", "ben_Beng"), ("Kashmiri (Devanagari)", "kas_Deva"), ("Sanskrit", "san_Deva"),
+    ("Bodo", "brx_Deva"), ("Maithili", "mai_Deva"), ("Santali", "sat_Olck"),
+    ("Dogri", "doi_Deva"), ("Malayalam", "mal_Mlym"), ("Sindhi (Arabic)", "snd_Arab"),
+    ("English", "eng_Latn"), ("Marathi", "mar_Deva"), ("Sindhi (Devanagari)", "snd_Deva"),
+    ("Konkani", "gom_Deva"), ("Manipuri (Bengali)", "mni_Beng"), ("Tamil", "tam_Taml"),
+    ("Gujarati", "guj_Gujr"), ("Manipuri (Meitei)", "mni_Mtei"), ("Telugu", "tel_Telu"),
+    ("Hindi", "hin_Deva"), ("Nepali", "npi_Deva"), ("Urdu", "urd_Arab"),
+    ("Kannada", "kan_Knda"), ("Odia", "ory_Orya")
+]
 
-# Create a Gradio interface
+# Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("### Indic Translations")
-    input_text = gr.Textbox(label="Input Text", placeholder="Enter text to translate")
-    src_lang = gr.Dropdown(label="Source Language", choices=list(LANGUAGES.keys()))
-    tgt_lang = gr.Dropdown(label="Target Language", choices=list(LANGUAGES.keys()))
-    translate_button = gr.Button("Translate")
-    translation_output = gr.Textbox(label="Translation", interactive=False)
+    gr.Markdown("# IndicTrans2 Translation")
+    with gr.Row():
+        with gr.Column():
+            input_text = gr.Textbox(label="Input Text")
+            src_lang = gr.Dropdown(label="Source Language", choices=[lang[0] for lang in languages], type="value")
+            tgt_lang = gr.Dropdown(label="Target Language", choices=[lang[0] for lang in languages], type="value")
+            translate_button = gr.Button("Translate")
+    
+    output_text = gr.Textbox(label="Translated Output")
 
-    @translate_button.click
-    def on_translate(text, src_lang, tgt_lang):
-        translation = translate(text, LANGUAGES[src_lang], LANGUAGES[tgt_lang])
-        translation_output.value = translation
+    # Call translate function when button is clicked
+    translate_button.click(fn=translate, inputs=[input_text, src_lang, tgt_lang], outputs=output_text)
 
 demo.launch()