Update app.py

app.py

@@ -1,39 +1,33 @@
-import requests
 import gradio as gr
-from dotenv import load_dotenv
-import os
-
-# Load environment variables from .env file
-load_dotenv()
-HF_TOKEN = os.getenv("HF_TOKEN")
-
-model_name = "Helsinki-NLP/opus-mt-en-fr"
-API_URL = f"https://api-inference.huggingface.co/models/{model_name}"
-headers = {"Authorization": f"Bearer {HF_TOKEN}"}
-
-
-def query(payload):
-    # HTTP POST Request
-    response = requests.post(API_URL, headers=headers, json=payload)
-    return response.json()
-
-
-def translate(input_text):
-    # API Request:
-    response = query({
-        "inputs": input_text,
-        "options": {"wait_for_model": True}
-    })
-
-    translation = response[0]["translation_text"]
-
-    return translation
-
-
-translator = gr.Interface(fn=translate,
-                          inputs=[gr.Textbox(label="Input Text", placeholder="Input Text To Be Translated")],
-                          outputs=gr.Textbox(label="Translation"),
-                          title="tTranslatorR-Opus"
-                          )
-
-translator.launch()
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+import torch
+
+# this model was loaded from https://hf.co/models
+model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
+tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
+device = 0 if torch.cuda.is_available() else -1
+LANGS = ["ace_Arab", "eng_Latn", "fra_Latn", "spa_Latn"]
+
+def translate(text, src_lang, tgt_lang):
+    """
+    Translate the text from source lang to target lang
+    """
+    translation_pipeline = pipeline("translation", model=model, tokenizer=tokenizer, src_lang=src_lang, tgt_lang=tgt_lang, max_length=400, device=device)
+    result = translation_pipeline(text)
+    return result[0]['translation_text']
+
+demo = gr.Interface(
+    fn=translate,
+    inputs=[
+        gr.components.Textbox(label="Text"),
+        gr.components.Dropdown(label="Source Language", choices=LANGS),
+        gr.components.Dropdown(label="Target Language", choices=LANGS),
+    ],
+    outputs=["text"],
+    examples=[["Building a translation demo with Gradio is so easy!", "eng_Latn", "spa_Latn"]],
+    cache_examples=False,
+    title="Translation Demo",
+    description="This demo is a simplified version of the original [NLLB-Translator](https://huggingface.co/spaces/Narrativaai/NLLB-Translator) space"
+)
+
+demo.launch()