Add per-word confidence, global confidence and uncertainty

app.py

@@ -1,12 +1,13 @@
 import io
 
 import torch
+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
 
-app = FastAPI(title="MMS Speech-to-Text API", version="1.0.0")
+app = FastAPI(title="MMS Speech-to-Text API", version="2.0.0")
 
 MODEL_ID = "facebook/mms-1b-all"
 processor = None
@@ -48,9 +49,56 @@ async def transcribe(file: UploadFile = File(...)):
     inputs = processor(audio, sampling_rate=sampling_rate, return_tensors="pt")
 
     with torch.no_grad():
-        logits = model(**inputs).logits
+        outputs = model(**inputs)
+        logits = outputs.logits  # (batch, time, vocab)
+
+    probs = F.softmax(logits, dim=-1)
+
+    predicted_ids = torch.argmax(probs, dim=-1)[0]
+    token_probs = torch.max(probs, dim=-1).values[0]
 
-    predicted_ids = torch.argmax(logits, dim=-1)[0]
     transcription = processor.decode(predicted_ids)
 
-    return JSONResponse({"transcription": transcription})
+    tokens = processor.tokenizer.convert_ids_to_tokens(predicted_ids)
+
+    words = []
+    current_word = ""
+    current_confs = []
+    prev_token = None
+
+    for tok, conf in zip(tokens, token_probs):
+        if tok == "<pad>":
+            continue
+        if tok == prev_token:
+            continue
+        prev_token = tok
+
+        if tok == "|":
+            if current_word:
+                words.append({
+                    "word": current_word,
+                    "confidence": float(sum(current_confs) / len(current_confs))
+                })
+            current_word = ""
+            current_confs = []
+        else:
+            current_word += tok
+            current_confs.append(conf.item())
+
+    if current_word:
+        words.append({
+            "word": current_word,
+            "confidence": float(sum(current_confs) / len(current_confs))
+        })
+
+    global_conf = float(token_probs.mean().item())
+
+    entropy = -torch.sum(probs * torch.log(probs + 1e-10), dim=-1)
+    uncertainty = float(entropy.mean().item())
+
+    return JSONResponse({
+        "transcription": transcription,
+        "confidence": global_conf,
+        "uncertainty": uncertainty,
+        "words": words
+    })