Update app.py

app.py

@@ -1,71 +1,44 @@
 import gradio as gr
 from transformers import pipeline, MarianTokenizer, AutoModelForSeq2SeqLM
 import torch
-import unicodedata
-import re
-import whisper
 import tempfile
 import os
-
-import nltk
-nltk.download('punkt')
-from nltk.tokenize import sent_tokenize
-
+import whisper
 import fitz  # PyMuPDF
 import docx
 from bs4 import BeautifulSoup
 import markdown2
 import chardet
+import re
 
 # Device setup
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-# Model configuration
-MODELS = {
-    "english_wolof": {
-        "model_name": "LocaleNLP/localenlp-eng-wol-0.03",
-        "target_tag": ">>wol<<"
-    },
-    "wolof_english": {
-        "model_name": "LocaleNLP/localenlp-wol-eng-0.03",
-        "target_tag": ">>eng<<"
-    },
-    "english_hausa": {
-        "model_name": "LocaleNLP/localenlp-eng-hau-0.01",
-        "target_tag": ">>hau<<"
-    },
-    "hausa_english": {
-        "model_name": "LocaleNLP/localenlp-hau-eng-0.01",
-        "target_tag": ">>eng<<"
-    }
-}
-
-# Global variables
+# Global model holders
 translator = None
-current_model = None
 whisper_model = None
 
+# Model configurations
+MODELS = {
+    ("English", "Wolof"): {"model_name": "LocaleNLP/localenlp-eng-wol-0.03", "tag": ">>wol<<"},
+    ("Wolof", "English"): {"model_name": "LocaleNLP/localenlp-wol-eng-0.03", "tag": ">>eng<<"},
+    ("English", "Hausa"): {"model_name": "LocaleNLP/localenlp-eng-hau-0.01", "tag": ">>hau<<"},
+    ("Hausa", "English"): {"model_name": "LocaleNLP/localenlp-hau-eng-0.01", "tag": ">>eng<<"},
+}
+
 HF_TOKEN = os.getenv("HF_TOKEN")
 
-def load_translation_model(input_lang, output_lang):
-    global translator, current_model
-    
-    model_key = f"{input_lang.lower()}_{output_lang.lower()}"
-    if model_key not in MODELS:
-        raise ValueError(f"Translation from {input_lang} to {output_lang} is not supported")
-    
-    if current_model != model_key or translator is None:
-        model_config = MODELS[model_key]
-        model = AutoModelForSeq2SeqLM.from_pretrained(model_config["model_name"], token=HF_TOKEN).to(device)
-        tokenizer = MarianTokenizer.from_pretrained(model_config["model_name"], token=HF_TOKEN)
-        translator = {
-            "pipeline": pipeline("translation", model=model, tokenizer=tokenizer, 
-                               device=0 if device.type == 'cuda' else -1),
-            "target_tag": model_config["target_tag"]
-        }
-        current_model = model_key
-    
-    return translator
+def load_model(input_lang, output_lang):
+    global translator
+    key = (input_lang, output_lang)
+    if key not in MODELS:
+        raise ValueError("Language pair not supported.")
+    cfg = MODELS[key]
+    if translator is None or translator.model.config._name_or_path != cfg["model_name"]:
+        model = AutoModelForSeq2SeqLM.from_pretrained(cfg["model_name"], token=HF_TOKEN).to(device)
+        tokenizer = MarianTokenizer.from_pretrained(cfg["model_name"], token=HF_TOKEN)
+        translator = pipeline("translation", model=model, tokenizer=tokenizer, device=0 if device.type=='cuda' else -1)
+    return translator, cfg["tag"]
 
 def load_whisper_model():
     global whisper_model
@@ -100,22 +73,16 @@ def extract_text_from_file(uploaded_file):
         with fitz.open(stream=content, filetype="pdf") as doc:
             return "\n".join([page.get_text() for page in doc])
     elif file_type == "docx":
-        if isinstance(uploaded_file, str):
-            doc = docx.Document(file_path)
-        else:
-            doc = docx.Document(uploaded_file)
+        doc = docx.Document(file_path if isinstance(uploaded_file, str) else uploaded_file)
         return "\n".join([para.text for para in doc.paragraphs])
     else:
         encoding = chardet.detect(content)['encoding']
-        if encoding:
-            content = content.decode(encoding, errors='ignore')
+        content = content.decode(encoding, errors='ignore') if encoding else content
         if file_type in ("html", "htm"):
-            soup = BeautifulSoup(content, "html.parser")
-            return soup.get_text()
+            return BeautifulSoup(content, "html.parser").get_text()
         elif file_type == "md":
             html = markdown2.markdown(content)
-            soup = BeautifulSoup(html, "html.parser")
-            return soup.get_text()
+            return BeautifulSoup(html, "html.parser").get_text()
         elif file_type == "srt":
             return re.sub(r"\d+\n\d{2}:\d{2}:\d{2},\d{3} --> .*?\n", "", content)
         elif file_type in ("txt", "text"):
@@ -123,11 +90,8 @@ def extract_text_from_file(uploaded_file):
         else:
             raise ValueError("Unsupported file type")
 
-def translate(text, input_lang, output_lang):
-    translator = load_translation_model(input_lang, output_lang)
-    lang_tag = translator["target_tag"]
-    translation_pipeline = translator["pipeline"]
-
+def translate_text(text, input_lang, output_lang):
+    translator, tag = load_model(input_lang, output_lang)
     paragraphs = text.split("\n")
     translated_output = []
 
@@ -137,9 +101,8 @@ def translate(text, input_lang, output_lang):
                 translated_output.append("")
                 continue
             sentences = [s.strip() for s in para.split('. ') if s.strip()]
-            formatted = [f"{lang_tag} {s}" for s in sentences]
-
-            results = translation_pipeline(formatted,
+            formatted = [f"{tag} {s}" for s in sentences]
+            results = translator(formatted,
                                  max_length=5000,
                                  num_beams=5,
                                  early_stopping=True,
@@ -148,43 +111,28 @@ def translate(text, input_lang, output_lang):
                                  length_penalty=1.2)
             translated_sentences = [r['translation_text'].capitalize() for r in results]
             translated_output.append('. '.join(translated_sentences))
-
     return "\n".join(translated_output)
 
-def process_input(input_mode, text, audio_file, file_obj, input_lang):
-    input_text = ""
+def process_input(input_mode, input_lang, text, audio_file, file_obj):
+    if input_mode == "Audio" and input_lang != "English":
+        raise ValueError("Audio input must be in English.")
     if input_mode == "Text":
-        input_text = text
-    elif input_mode == "Audio":
-        if audio_file is not None:
-            input_text = transcribe_audio(audio_file)
-    elif input_mode == "File":
-        if file_obj is not None:
-            input_text = extract_text_from_file(file_obj)
-    return input_text
-
-def translate_and_return(text, input_lang, output_lang):
-    if not text.strip():
-        return "No input text to translate."
-    return translate(text, input_lang, output_lang)
-
-def update_input_lang_dropdown(input_mode):
-    if input_mode == "Audio":
-        return gr.Dropdown(value="English", interactive=False)
-    else:
-        return gr.Dropdown(interactive=True)
-
-# Gradio UI components
+        return text
+    elif input_mode == "Audio" and audio_file is not None:
+        return transcribe_audio(audio_file)
+    elif input_mode == "File" and file_obj is not None:
+        return extract_text_from_file(file_obj)
+    return ""
+
+# Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## LocaleNLP Translator")
-    gr.Markdown("Translate between English, Wolof, and Hausa using Localenlp models.")
+    gr.Markdown("## LocaleNLP Multi-language Translator")
+    gr.Markdown("Translate between English, Wolof, and Hausa. Audio input only accepts English.")
 
     with gr.Row():
-        input_mode = gr.Radio(choices=["Text", "Audio", "File"], label="Select input mode", value="Text")
-    
-    with gr.Row():
-        input_lang = gr.Dropdown(choices=["English", "Wolof", "Hausa"], label="Input Language", value="English")
-        output_lang = gr.Dropdown(choices=["English", "Wolof", "Hausa"], label="Output Language", value="Hausa")
+        input_mode = gr.Radio(choices=["Text", "Audio", "File"], label="Input type", value="Text")
+        input_lang = gr.Dropdown(choices=["English", "Wolof", "Hausa"], label="Input language", value="English")
+        output_lang = gr.Dropdown(choices=["English", "Wolof", "Hausa"], label="Output language", value="Hausa")
 
     input_text = gr.Textbox(label="Enter text", lines=10, visible=True)
     audio_input = gr.Audio(label="Upload audio (.wav, .mp3, .m4a)", type="filepath", visible=False)
@@ -202,22 +150,20 @@ with gr.Blocks() as demo:
             extracted_text: gr.update(value="", visible=True),
             output_text: gr.update(value="")
         }
-
     input_mode.change(fn=update_visibility, inputs=input_mode, outputs=[input_text, audio_input, file_input, extracted_text, output_text])
-    input_mode.change(fn=update_input_lang_dropdown, inputs=input_mode, outputs=input_lang)
 
-    def handle_process(mode, text, audio, file_obj, in_lang):
+    def handle_process(mode, lang_in, text, audio, file_obj):
         try:
-            extracted = process_input(mode, text, audio, file_obj, in_lang)
+            extracted = process_input(mode, lang_in, text, audio, file_obj)
             return extracted, ""
         except Exception as e:
             return "", f"Error: {str(e)}"
+    translate_button.click(fn=handle_process, inputs=[input_mode, input_lang, input_text, audio_input, file_input], outputs=[extracted_text, output_text])
 
-    translate_button.click(fn=handle_process, inputs=[input_mode, input_text, audio_input, file_input, input_lang], outputs=[extracted_text, output_text])
-
-    def handle_translate(text, in_lang, out_lang):
-        return translate_and_return(text, in_lang, out_lang)
-
+    def handle_translate(text, lang_in, lang_out):
+        if not text.strip():
+            return "No input text to translate."
+        return translate_text(text, lang_in, lang_out)
     translate_button.click(fn=handle_translate, inputs=[extracted_text, input_lang, output_lang], outputs=output_text)
 
-demo.launch()
+demo.launch()