Add SeamlessM4T v2 Large STT + NLLB-200

app.py

@@ -1,18 +1,36 @@
 import gradio as gr
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from transformers import AutoProcessor, SeamlessM4Tv2ForSpeechToText, AutoModelForSeq2SeqLM, AutoTokenizer
 import torch
+import numpy as np
+
+# ============================================================
+# 🚀 Load Models
+# ============================================================
 
-# Load NLLB-200 (distilled for speed)
-MODEL_NAME = "facebook/nllb-200-distilled-600M"
-print(f"Loading {MODEL_NAME}...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model = model.to(device)
-print(f"Model loaded on {device}")
+print(f"🖥️ Device: {device}")
+
+# SeamlessM4T v2 Large for STT
+print("📥 Loading SeamlessM4T v2 Large...")
+stt_model_name = "facebook/seamless-m4t-v2-large"
+stt_processor = AutoProcessor.from_pretrained(stt_model_name)
+stt_model = SeamlessM4Tv2ForSpeechToText.from_pretrained(stt_model_name).to(device)
+print("✅ SeamlessM4T v2 Large loaded")
+
+# NLLB-200 for Translation
+print("📥 Loading NLLB-200...")
+nllb_model_name = "facebook/nllb-200-distilled-600M"
+nllb_tokenizer = AutoTokenizer.from_pretrained(nllb_model_name)
+nllb_model = AutoModelForSeq2SeqLM.from_pretrained(nllb_model_name).to(device)
+print("✅ NLLB-200 loaded")
 
-# Language codes
-LANGS = {
+print("🎉 All models ready!")
+
+# ============================================================
+# Language Codes
+# ============================================================
+
+NLLB_LANGS = {
     "English": "eng_Latn",
     "French": "fra_Latn",
     "Arabic": "arb_Arab",
@@ -30,33 +48,105 @@ LANGS = {
     "Hindi": "hin_Deva",
 }
 
+STT_LANGS = {
+    "English": "eng",
+    "French": "fra",
+    "Arabic": "arb",
+    "Spanish": "spa",
+    "German": "deu",
+    "Italian": "ita",
+    "Portuguese": "por",
+    "Chinese": "cmn",
+    "Japanese": "jpn",
+    "Korean": "kor",
+    "Russian": "rus",
+    "Turkish": "tur",
+    "Dutch": "nld",
+    "Hindi": "hin",
+}
+
+# ============================================================
+# STT Function (SeamlessM4T v2 Large)
+# ============================================================
+
+def stt(audio, src_lang):
+    """Speech-to-Text using SeamlessM4T v2 Large"""
+    if audio is None:
+        return ""
+    
+    # Handle tuple input from Gradio
+    if isinstance(audio, tuple):
+        sample_rate, audio_data = audio
+        audio_data = audio_data.astype(np.float32)
+        if audio_data.max() > 1.0:
+            audio_data = audio_data / 32768.0
+    else:
+        return "Error: Invalid audio format"
+    
+    src_code = STT_LANGS.get(src_lang, "eng")
+    
+    inputs = stt_processor(
+        audios=audio_data,
+        sampling_rate=sample_rate,
+        return_tensors="pt"
+    ).to(device)
+    
+    with torch.no_grad():
+        output_tokens = stt_model.generate(
+            **inputs,
+            tgt_lang=src_code,
+            generate_speech=False
+        )
+    
+    text = stt_processor.decode(output_tokens[0], skip_special_tokens=True)
+    return text
+
+# ============================================================
+# Translation Function (NLLB-200)
+# ============================================================
+
 def translate(text, src_lang, tgt_lang):
+    """Translation using NLLB-200"""
     if not text.strip():
         return ""
     
-    src_code = LANGS.get(src_lang, "eng_Latn")
-    tgt_code = LANGS.get(tgt_lang, "fra_Latn")
+    src_code = NLLB_LANGS.get(src_lang, "eng_Latn")
+    tgt_code = NLLB_LANGS.get(tgt_lang, "fra_Latn")
     
-    tokenizer.src_lang = src_code
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512).to(device)
+    nllb_tokenizer.src_lang = src_code
+    inputs = nllb_tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512).to(device)
     
-    forced_bos_token_id = tokenizer.convert_tokens_to_ids(tgt_code)
+    forced_bos_token_id = nllb_tokenizer.convert_tokens_to_ids(tgt_code)
     
     with torch.no_grad():
-        outputs = model.generate(**inputs, forced_bos_token_id=forced_bos_token_id, max_length=512, num_beams=5)
-    
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-demo = gr.Interface(
-    fn=translate,
-    inputs=[
-        gr.Textbox(label="Text to translate", lines=3),
-        gr.Dropdown(list(LANGS.keys()), label="Source Language", value="English"),
-        gr.Dropdown(list(LANGS.keys()), label="Target Language", value="French"),
-    ],
-    outputs=gr.Textbox(label="Translation", lines=3),
-    title="NLLB-200 Translation API",
-    description="200 languages including Moroccan Arabic!",
-)
+        outputs = nllb_model.generate(**inputs, forced_bos_token_id=forced_bos_token_id, max_length=512, num_beams=5)
+    
+    return nllb_tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+# ============================================================
+# Gradio Interface
+# ============================================================
+
+with gr.Blocks(title="STTR - Speech & Translation API") as demo:
+    gr.Markdown("# 🌍 STTR - Speech-to-Text & Translation API")
+    gr.Markdown("**SeamlessM4T v2 Large** for STT + **NLLB-200** for Translation")
+    
+    with gr.Tab("🎤 STT (Speech-to-Text)"):
+        with gr.Row():
+            stt_audio = gr.Audio(label="Record/Upload Audio", type="numpy")
+            stt_lang = gr.Dropdown(list(STT_LANGS.keys()), label="Language", value="English")
+        stt_output = gr.Textbox(label="Transcription", lines=3)
+        stt_btn = gr.Button("🎤 Transcribe", variant="primary")
+        stt_btn.click(stt, inputs=[stt_audio, stt_lang], outputs=stt_output, api_name="stt")
+    
+    with gr.Tab("🌍 Translation"):
+        with gr.Row():
+            trans_text = gr.Textbox(label="Text to translate", lines=3)
+        with gr.Row():
+            trans_src = gr.Dropdown(list(NLLB_LANGS.keys()), label="Source", value="English")
+            trans_tgt = gr.Dropdown(list(NLLB_LANGS.keys()), label="Target", value="French")
+        trans_output = gr.Textbox(label="Translation", lines=3)
+        trans_btn = gr.Button("🌍 Translate", variant="primary")
+        trans_btn.click(translate, inputs=[trans_text, trans_src, trans_tgt], outputs=trans_output, api_name="translate")
 
 demo.launch()

requirements.txt

@@ -1,5 +1,8 @@
-transformers>=4.30.0
+transformers>=4.40.0
 torch>=2.0.0
 sentencepiece
 protobuf
-gradio
+gradio>=4.0.0
+numpy
+scipy
+torchaudio