Upload app.py

nllb gradio model

app.py

@@ -0,0 +1,62 @@
+import gradio as gr
+import torch
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+
+# Load model and tokenizer
+model_name = "./model"
+tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+
+# Check if CUDA (GPU) is available, otherwise fall back to CPU
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
+
+# Define lang code to token ID mapping manually
+lang_code_to_token = {
+    "eng_Latn": "__eng_Latn__",
+    "deu_Latn": "__deu_Latn__",
+    "fra_Latn": "__fra_Latn__",
+    "spa_Latn": "__spa_Latn__"
+}
+
+# Choose target language (change as needed)
+target_lang_code = "eng_Latn"  # ← change to "deu_Latn" for German
+src_lang_code = "mya"          # Burmese
+
+def translate(text):
+    if not text.strip():
+        return "Please enter some Burmese text."
+
+    # Set source language for tokenizer
+    tokenizer.src_lang = src_lang_code
+
+    # Tokenize input
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to(device)
+
+    # ✅ Get forced BOS token ID manually
+    lang_token = lang_code_to_token[target_lang_code]
+    forced_bos_token_id = tokenizer.convert_tokens_to_ids(lang_token)
+
+    # Generate translation
+    with torch.no_grad():
+        generated_tokens = model.generate(
+            **inputs,
+            forced_bos_token_id=forced_bos_token_id,
+            max_length=128
+        )
+
+    return tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
+
+# Gradio UI
+with gr.Blocks(title="Burmese to English Translator") as demo:
+    gr.Markdown("## 🇲🇲 ➡️ 🇬🇧 Burmese to English Translator")
+    with gr.Row():
+        with gr.Column():
+            input_text = gr.Textbox(label="Burmese Text", placeholder="Type Burmese here...", lines=4)
+            btn = gr.Button("Translate")
+        with gr.Column():
+            output_text = gr.Textbox(label="Translation", interactive=False)
+
+    btn.click(translate, inputs=input_text, outputs=output_text)
+
+demo.launch()