Update app.py

app.py

@@ -1,85 +1,64 @@
-import os
 import gradio as gr
-from transformers import pipeline
 from tts_engine import TTSEngine
+from translation import Translator, CustomTranslator
+from data_manager import save_uploaded_file, convert_to_jsonl
+from training.train_translation import train_from_jsonl
+from stt_engine import STTEngine
+import os
 
-# STT: Whisper
-stt_pipeline = pipeline("automatic-speech-recognition", model="openai/whisper-small", device=-1)
-
-# Translation: MarianMT (generic English <-> multilingual)
-translator_pipeline = pipeline("translation", model="Helsinki-NLP/opus-mt-mul-en")
-
+# Init engines
+stt_engine = STTEngine()
 tts_engine = TTSEngine(use_coqui=True)
+translator = CustomTranslator() if os.path.exists("./training/outputs/model") else Translator()
 
-LANGUAGES = [
-    "english", "yoruba", "igbo", "hausa", "pidgin",
-    "esan", "tiv", "calabar", "benin"
-]
+LANGUAGES = ["english", "yoruba", "igbo", "hausa", "pidgin", "esan", "tiv", "calabar", "benin"]
 
-def transcribe(audio, language="en"):
-    """Speech to text using Whisper."""
+def handle_conversation(audio, src_lang, tgt_lang, clone_voice):
     if audio is None:
-        return ""
-    result = stt_pipeline(audio, generate_kwargs={"language": language})
-    return result["text"]
-
-def translate(text, src_lang, tgt_lang):
-    """Dummy translation with Marian (you can extend with Nigerian mappings)."""
-    if not text:
-        return ""
-    # For Nigerian languages not covered by Marian, just return the same text
-    if src_lang not in ["english", "en"] or tgt_lang not in ["english", "en"]:
-        return f"[{src_lang}->{tgt_lang}] {text}"
-    translated = translator_pipeline(text)[0]["translation_text"]
-    return translated
+        return "", None
 
-def handle_conversation(audio, src_lang, tgt_lang, clone_voice):
-    """One side speaks -> STT -> Translate -> TTS"""
-    # Step 1: STT
-    text = transcribe(audio, language="en" if src_lang == "english" else None)
+    # Step 1: Speech to Text
+    text = stt_engine.transcribe(audio, language=src_lang)
 
     # Step 2: Translate
-    translated = translate(text, src_lang, tgt_lang)
+    translated = translator.translate(text, src_lang, tgt_lang)
 
-    # Step 3: TTS
+    # Step 3: Text to Speech
     audio_path = tts_engine.speak(translated, lang=tgt_lang, voice_clone=clone_voice)
 
     return translated, audio_path
 
+def admin_upload(file):
+    file_path = save_uploaded_file(file, file.name)
+    jsonl_path = convert_to_jsonl(file_path)
+    train_from_jsonl(jsonl_path)
+    return "✅ Training done. Model updated!"
 
-with gr.Blocks(title="🌍 Two-Way Translation Assistant") as demo:
+with gr.Blocks(title="🌍 Two-Way Voice Translator") as demo:
     gr.Markdown("# 🌍 Nigerian Two-Way Voice Translator")
-    gr.Markdown("Speak in your language, hear it in theirs. Supports English ↔ Nigerian languages.")
-
-    with gr.Row():
-        with gr.Column():
-            gr.Markdown("### 🧑 Speaker A")
+    with gr.Tab("Translator"):
+        with gr.Row():
             src_lang = gr.Dropdown(LANGUAGES, value="english", label="Speaker A Language")
-            audio_in_a = gr.Audio(sources=["microphone"], type="filepath", label="🎤 Speak here")
-            translated_a = gr.Textbox(label="Translated Text", interactive=False)
-            audio_out_a = gr.Audio(label="🔊 Translation Audio")
-
-        with gr.Column():
-            gr.Markdown("### 👩 Speaker B")
-            tgt_lang = gr.Dropdown(LANGUAGES, value="yoruba", label="Speaker B Language")
-            audio_in_b = gr.Audio(sources=["microphone"], type="filepath", label="🎤 Reply here")
-            translated_b = gr.Textbox(label="Translated Text", interactive=False)
-            audio_out_b = gr.Audio(label="🔊 Translation Audio")
-
-    clone_voice = gr.Checkbox(value=False, label="🎙️ Use my cloned voice (if my_voice.wav exists)")
-
-    # Wire up A -> B
-    audio_in_a.change(
-        handle_conversation,
-        inputs=[audio_in_a, src_lang, tgt_lang, clone_voice],
-        outputs=[translated_a, audio_out_a]
-    )
-
-    # Wire up B -> A
-    audio_in_b.change(
-        handle_conversation,
-        inputs=[audio_in_b, tgt_lang, src_lang, clone_voice],
-        outputs=[translated_b, audio_out_b]
-    )
+            tgt_lang = gr.Dropdown(LANGUAGES, value="hausa", label="Speaker B Language")
+
+        with gr.Row():
+            audio_in = gr.Audio(sources=["microphone"], type="filepath", label="🎤 Speak")
+            translated = gr.Textbox(label="Translated Text", interactive=False)
+            audio_out = gr.Audio(label="🔊 Translation Audio")
+
+        clone_voice = gr.Checkbox(value=False, label="🎙️ Use my cloned voice (if my_voice.wav exists)")
+
+        audio_in.change(
+            handle_conversation,
+            inputs=[audio_in, src_lang, tgt_lang, clone_voice],
+            outputs=[translated, audio_out]
+        )
+
+    with gr.Tab("Admin (Training)"):
+        gr.Markdown("Upload Hausa ↔ English data (.csv, .xlsx, .tsv, .jsonl)")
+        file_in = gr.File(label="Upload dataset")
+        train_btn = gr.Button("🚀 Train Model")
+        output_box = gr.Textbox(label="Training Status")
+        train_btn.click(admin_upload, inputs=file_in, outputs=output_box)
 
 demo.launch()