Add LID endpoint using facebook/mms-lid-256

app.py

@@ -5,23 +5,31 @@ import torch.nn.functional as F
 import librosa
 from fastapi import FastAPI, File, UploadFile, HTTPException
 from fastapi.responses import JSONResponse
-from transformers import Wav2Vec2ForCTC, AutoProcessor
+from transformers import Wav2Vec2ForCTC, AutoProcessor, AutoFeatureExtractor, AutoModelForAudioClassification
 
 app = FastAPI(title="MMS Speech-to-Text API", version="2.0.0")
 
 MODEL_ID = "facebook/mms-1b-all"
+LID_MODEL_ID = "facebook/mms-lid-256"
+
 processor = None
 model = None
+lid_extractor = None
+lid_model = None
 
 
 @app.on_event("startup")
 async def load_model():
-    global processor, model
-    print("Loading MMS model...")
+    global processor, model, lid_extractor, lid_model
+    print("Loading MMS ASR model...")
     processor = AutoProcessor.from_pretrained(MODEL_ID)
     model = Wav2Vec2ForCTC.from_pretrained(MODEL_ID)
     model.eval()
-    print("Model loaded.")
+    print("Loading MMS LID model...")
+    lid_extractor = AutoFeatureExtractor.from_pretrained(LID_MODEL_ID)
+    lid_model = AutoModelForAudioClassification.from_pretrained(LID_MODEL_ID)
+    lid_model.eval()
+    print("All models loaded.")
 
 
 @app.get("/")
@@ -31,7 +39,11 @@ def root():
 
 @app.get("/health")
 def health():
-    return {"status": "ok", "model_loaded": model is not None}
+    return {
+        "status": "ok",
+        "asr_model_loaded": model is not None,
+        "lid_model_loaded": lid_model is not None,
+    }
 
 
 @app.post("/transcribe")
@@ -50,7 +62,7 @@ async def transcribe(file: UploadFile = File(...)):
 
     with torch.no_grad():
         outputs = model(**inputs)
-        logits = outputs.logits  # (batch, time, vocab)
+        logits = outputs.logits
 
     probs = F.softmax(logits, dim=-1)
 
@@ -58,7 +70,6 @@ async def transcribe(file: UploadFile = File(...)):
     token_probs = torch.max(probs, dim=-1).values[0]
 
     transcription = processor.decode(predicted_ids)
-
     tokens = processor.tokenizer.convert_ids_to_tokens(predicted_ids)
 
     words = []
@@ -92,7 +103,6 @@ async def transcribe(file: UploadFile = File(...)):
         })
 
     global_conf = float(token_probs.mean().item())
-
     entropy = -torch.sum(probs * torch.log(probs + 1e-10), dim=-1)
     uncertainty = float(entropy.mean().item())
 
@@ -102,3 +112,27 @@ async def transcribe(file: UploadFile = File(...)):
         "uncertainty": uncertainty,
         "words": words
     })
+
+
+@app.post("/lid")
+async def language_identification(file: UploadFile = File(...)):
+    if lid_model is None or lid_extractor is None:
+        raise HTTPException(status_code=503, detail="LID model not loaded yet")
+
+    audio_bytes = await file.read()
+
+    try:
+        audio_input, _ = librosa.load(io.BytesIO(audio_bytes), sr=16000, mono=True)
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Failed to load audio: {e}")
+
+    inputs = lid_extractor(audio_input, sampling_rate=16000, return_tensors="pt")
+
+    with torch.no_grad():
+        outputs = lid_model(**inputs)
+        logits = outputs.logits
+
+    predicted_id = torch.argmax(logits, dim=-1).item()
+    predicted_lang = lid_model.config.id2label[predicted_id]
+
+    return JSONResponse({"language": predicted_lang})