Update app.py

app.py

@@ -1,9 +1,4 @@
-import os
-import os
-# Redirect Streamlit config/cache to a writable temp folder
-os.environ["STREAMLIT_HOME"] = "/tmp/.streamlit"
-
-import streamlit as st
+import gradio as gr
 from transformers import pipeline, MarianTokenizer, AutoModelForSeq2SeqLM
 import torch
 import unicodedata
@@ -16,31 +11,47 @@ import nltk
 nltk.download('punkt')
 from nltk.tokenize import sent_tokenize
 
-import fitz  # PyMuPDF for PDF
+import fitz  # PyMuPDF
 import docx
 from bs4 import BeautifulSoup
 import markdown2
 import chardet
 
-# --- Device selection ---
+# Device setup
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-# --- Load Wolof MarianMT model from HF hub ---
-@st.cache_resource
+# Load Wolof MarianMT model from HF hub (cached manually)
+translator = None
+whisper_model = None
+
 def load_wolof_model():
-    model_name = "LocaleNLP/eng_wolof"
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device)
-    tokenizer = MarianTokenizer.from_pretrained(model_name)
-    translator = pipeline("translation", model=model, tokenizer=tokenizer, device=0 if device.type == 'cuda' else -1)
+    global translator
+    if translator is None:
+        model_name = "LocaleNLP/eng_wolof"
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device)
+        tokenizer = MarianTokenizer.from_pretrained(model_name)
+        translator = pipeline("translation", model=model, tokenizer=tokenizer, device=0 if device.type == 'cuda' else -1)
     return translator
 
-@st.cache_resource
 def load_whisper_model():
-    return whisper.load_model("base")
-
-def transcribe_audio(audio_path):
-    whisper_model = load_whisper_model()
-    return whisper_model.transcribe(audio_path)["text"]
+    global whisper_model
+    if whisper_model is None:
+        whisper_model = whisper.load_model("base")
+    return whisper_model
+
+def transcribe_audio(audio_file):
+    model = load_whisper_model()
+    # Save temp file if not a path
+    if isinstance(audio_file, str):
+        audio_path = audio_file
+    else:
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
+            tmp.write(audio_file.read())
+            audio_path = tmp.name
+    result = model.transcribe(audio_path)
+    if not isinstance(audio_file, str):
+        os.remove(audio_path)
+    return result["text"]
 
 def translate(text):
     translator = load_wolof_model()
@@ -97,73 +108,62 @@ def extract_text_from_file(uploaded_file):
         else:
             raise ValueError("Unsupported file type")
 
-def main():
-    st.set_page_config(page_title="LocaleNLP English-to-Wolof Translator", layout="wide", initial_sidebar_state="expanded")
-
-    with st.sidebar:
-        st.image("localenpl5.jpeg", use_container_width=True)
-        st.markdown("""
-        <h3 style='text-align: left; color: #4B8BBE;'>🌐 Wolof Translation</h3>
-        This app translates English text to Wolof (Senegal) using a custom MarianMT model hosted on Hugging Face.
-        """, unsafe_allow_html=True)
-
-    st.markdown("<h4 style='text-align: center; color: #306998;'>Translate English to Wolof</h4>", unsafe_allow_html=True)
-
-    col1, col2 = st.columns(2)
-    with col1:
-        input_mode = st.selectbox("Select input mode:", ("Text", "Audio", "File"))
-        st.markdown("<hr>", unsafe_allow_html=True)
-    with col2:
-        target_lang = "Wolof (Senegal)"
-        st.markdown(f"<p><b>Target language:</b> {target_lang}</p>", unsafe_allow_html=True)
-        st.markdown("<hr>", unsafe_allow_html=True)
-
-    col3, col4 = st.columns(2)
-    with col3:
-        input_text = ""
-        if input_mode == "Text":
-            input_text = st.text_area("✏️ Enter English text:", height=250)
-        elif input_mode == "Audio":
-            audio_file = st.file_uploader("🔊 Upload audio (.wav, .mp3, .m4a)", type=["wav", "mp3", "m4a"])
-            if audio_file:
-                with tempfile.NamedTemporaryFile(delete=False, suffix=f".{audio_file.type.split('/')[-1]}") as tmp:
-                    tmp.write(audio_file.read())
-                    tmp_path = tmp.name
-                with st.spinner("Transcribing..."):
-                    input_text = transcribe_audio(tmp_path)
-                    os.remove(tmp_path)
-                    st.text_area("📝 Transcribed Text:", value=input_text, height=150)
-        elif input_mode == "File":
-            uploaded_file = st.file_uploader("📄 Upload document (PDF, Word, HTML, Markdown, SRT, TXT)", 
-                                             type=["pdf", "docx", "html", "htm", "md", "srt", "txt"])
-            if uploaded_file:
-                try:
-                    input_text = extract_text_from_file(uploaded_file)
-                    st.text_area("📃 Extracted Text:", value=input_text, height=200)
-                except Exception as e:
-                    st.error(f"Error extracting text: {str(e)}")
-
-    with col4:
-        if input_text:
-            with st.spinner("Translating..."):
-                translated_text = translate(input_text)
-            st.write(f"Output in {target_lang}")
-            st.success(translated_text)
-
-            st.download_button(
-                label="💾 Download Translation",
-                data=translated_text,
-                file_name=f"translated_{target_lang.replace(' ', '_').lower()}.txt",
-                mime="text/plain"
-            )
-        else:
-            st.info("Translation will appear here.")
-
-    st.markdown("""
-    <hr>
-    <div style='text-align: center; color: #4B8BBE; font-size: 0.9rem'>
-        LocaleNLP © 2025 • Empowering communities through AI and language
-    </div>""", unsafe_allow_html=True)
-
-if __name__ == "__main__":
-    main()
+def process_input(input_mode, text, audio_file, file_obj):
+    input_text = ""
+    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):
+    if not text.strip():
+        return "No input text to translate."
+    return translate(text)
+
+# Gradio UI components
+with gr.Blocks() as demo:
+    gr.Markdown("## LocaleNLP English-to-Wolof Translator")
+    gr.Markdown("Upload English text, audio, or document to translate to Wolof using a custom MarianMT model.")
+
+    with gr.Row():
+        input_mode = gr.Radio(choices=["Text", "Audio", "File"], label="Select input mode", value="Text")
+
+    input_text = gr.Textbox(label="Enter English text", lines=10, visible=True)
+    audio_input = gr.Audio(source="upload", type="file", label="Upload audio (.wav, .mp3, .m4a)", visible=False)
+    file_input = gr.File(file_types=['.pdf', '.docx', '.html', '.htm', '.md', '.srt', '.txt'], label="Upload document", visible=False)
+
+    extracted_text = gr.Textbox(label="Extracted / Transcribed Text", lines=10, interactive=False)
+    translate_button = gr.Button("Translate to Wolof")
+    output_text = gr.Textbox(label="Translated Wolof Text", lines=10, interactive=False)
+
+    def update_visibility(mode):
+        return {
+            input_text: gr.update(visible=(mode=="Text")),
+            audio_input: gr.update(visible=(mode=="Audio")),
+            file_input: gr.update(visible=(mode=="File")),
+            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])
+
+    def handle_process(mode, text, audio, file_obj):
+        try:
+            extracted = process_input(mode, 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_text, audio_input, file_input], outputs=[extracted_text, output_text])
+
+    def handle_translate(text):
+        return translate_and_return(text)
+
+    translate_button.click(fn=handle_translate, inputs=extracted_text, outputs=output_text)
+
+demo.launch()