Update app.py

app.py

@@ -2,20 +2,39 @@ from fastapi import FastAPI, UploadFile, Form
 from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
 import torch
 import torch.nn.functional as F
-import numpy as np
-from word_syllables import word_syllables
 import librosa
 
 app = FastAPI(title="Cognia Wav2Vec2 Speech API")
 
+# -------------------------------
+# MODEL SETUP
+# -------------------------------
 MODEL_DIR = "Artyomorax/cognia-wav2vec"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-SR = 16000
+SR = 16000  # sampling rate
 
 processor = Wav2Vec2Processor.from_pretrained(MODEL_DIR)
 model = Wav2Vec2ForCTC.from_pretrained(MODEL_DIR).to(DEVICE)
 model.eval()
 
+# -------------------------------
+# WORD SYLLABLES MAPPING
+# -------------------------------
+word_syllables = {
+    "otso": ["ot", "so"],
+    "ulap": ["u", "lap"],
+    "ubo": ["u", "bo"],
+    "anak": ["a", "nak"],
+    "aso": ["a", "so"],
+    "aklat": ["ak", "lat"],
+    "bahay": ["ba", "hay"],
+    "bata": ["ba", "ta"],
+    # ... continue filling your full dictionary ...
+}
+
+# -------------------------------
+# HELPERS
+# -------------------------------
 def get_embedding(syllables):
     token_ids = []
     for s in syllables:
@@ -27,6 +46,9 @@ def get_embedding(syllables):
     embeds = model.lm_head.weight[tokens]
     return embeds.mean(dim=0)
 
+# -------------------------------
+# API ENDPOINT
+# -------------------------------
 @app.post("/analyze")
 async def analyze_audio(file: UploadFile, target_word: str = Form(...)):
     # load audio
@@ -40,15 +62,18 @@ async def analyze_audio(file: UploadFile, target_word: str = Form(...)):
     transcription = processor.batch_decode(pred_ids, group_tokens=True)[0].lower().strip()
 
     # cosine similarity
-    target_sylls = word_syllables.get(target_word.lower(), [])
+    target_word_clean = target_word.lower().strip()
+    target_sylls = word_syllables.get(target_word_clean, [])
     pred_syllables = transcription.replace("-", " ").split()
+
     pred_emb = get_embedding(pred_syllables)
     ref_emb = get_embedding(target_sylls)
+
     cosine_sim = F.cosine_similarity(pred_emb.unsqueeze(0), ref_emb.unsqueeze(0)).item() if pred_emb is not None and ref_emb is not None else 0.0
 
     return {
         "transcription": transcription,
-        "target_word": target_word,
+        "target_word": target_word_clean,
         "cosine_similarity": round(float(cosine_sim), 2)
     }