Update app.py

app.py

@@ -1,9 +1,9 @@
 """
 VoiceAura Translation API
 Models:
-  1. SLPG/English_to_Urdu_Unsupervised_MT       (en → ur)
-  2. SLPG/Punjabi_Shahmukhi_to_Gurmukhi          (pa-s → pa-g)
-  3. SLPG/Punjabi_Gurmukhi_to_Shahmukhi          (pa-g → pa-s)
+  1. SLPG/English_to_Urdu_Unsupervised_MT            (en    → ur)
+  2. SLPG/Punjabi_Shahmukhi_to_Gurmukhi_Transliteration (pa-s → pa-g)
+  3. SLPG/Punjabi_Gurmukhi_to_Shahmukhi_Transliteration (pa-g → pa-s)
 """
 
 from fastapi import FastAPI
@@ -13,6 +13,11 @@ import os, requests, argparse, torch
 
 # ✅ PyTorch 2.6 fix
 torch.serialization.add_safe_globals([argparse.Namespace])
+_original_torch_load = torch.load
+
+def patched_torch_load(*args, **kwargs):
+    kwargs["weights_only"] = False
+    return _original_torch_load(*args, **kwargs)
 
 app = FastAPI()
 app.add_middleware(
@@ -22,7 +27,7 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# ── Model URLs ───────────────────────────────────────────
+# ── Model configs ────────────────────────────────────────
 MODELS_CONFIG = {
     "en-ur": {
         "files": {
@@ -32,6 +37,7 @@ MODELS_CONFIG = {
         },
         "dir":        "models/en_ur",
         "checkpoint": "checkpoint_8_96000.pt",
+        "detokenize": False,
         "instance":   None,
     },
     "pa-s-pa-g": {
@@ -42,6 +48,18 @@ MODELS_CONFIG = {
         },
         "dir":        "models/pa_s_pa_g",
         "checkpoint": "checkpoint_5_78000.pt",
+        "detokenize": True,
+        "instance":   None,
+    },
+    "pa-g-pa-s": {
+        "files": {
+            "checkpoint_13_129000.pt": "https://huggingface.co/SLPG/Punjabi_Gurmukhi_to_Shahmukhi_Transliteration/resolve/main/checkpoint_13_129000.pt",
+            "dict.pa.txt":             "https://huggingface.co/SLPG/Punjabi_Gurmukhi_to_Shahmukhi_Transliteration/resolve/main/dict.pa.txt",
+            "dict.pk.txt":             "https://huggingface.co/SLPG/Punjabi_Gurmukhi_to_Shahmukhi_Transliteration/resolve/main/dict.pk.txt",
+        },
+        "dir":        "models/pa_g_pa_s",
+        "checkpoint": "checkpoint_13_129000.pt",
+        "detokenize": True,
         "instance":   None,
     },
 }
@@ -60,22 +78,18 @@ def download_file(url: str, path: str):
                 f.write(chunk)
     print(f"[✓] Done: {path}")
 
-def patched_torch_load(*args, **kwargs):
-    kwargs["weights_only"] = False
-    return _original_torch_load(*args, **kwargs)
-
-_original_torch_load = torch.load
+def detokenize(text: str) -> str:
+    """Remove fairseq BPE tokens (▁ symbols)"""
+    return text.replace("▁", "").strip()
 
 def load_model(pair: str):
     cfg = MODELS_CONFIG[pair]
     if cfg["instance"] is not None:
         return cfg["instance"]
 
-    # Download files
     for fname, url in cfg["files"].items():
         download_file(url, os.path.join(cfg["dir"], fname))
 
-    # Patch torch.load for fairseq
     torch.load = patched_torch_load
     from fairseq.models.transformer import TransformerModel
     model = TransformerModel.from_pretrained(
@@ -89,13 +103,13 @@ def load_model(pair: str):
     print(f"[✓] Model ready: {pair}")
     return model
 
-# ── Startup — load all models ────────────────────────────
+# ── Startup ──────────────────────────────────────────────
 @app.on_event("startup")
 async def startup():
     for pair in MODELS_CONFIG:
         load_model(pair)
 
-# ── Endpoints ────────────────────────────────────────────
+# ── API ──────────────────────────────────────────────────
 class Req(BaseModel):
     text: str
     from_lang: str = "en"
@@ -104,7 +118,7 @@ class Req(BaseModel):
 @app.get("/")
 def root():
     loaded = {k: MODELS_CONFIG[k]["instance"] is not None for k in MODELS_CONFIG}
-    return {"status": "VoiceAura API ✓", "models": loaded}
+    return {"status": "VoiceAura API ✓", "models_loaded": loaded}
 
 @app.post("/translate")
 def translate(req: Req):
@@ -117,9 +131,15 @@ def translate(req: Req):
         return {"success": False, "translation": f"⚠️ Pair '{pair}' not supported."}
 
     try:
+        cfg    = MODELS_CONFIG[pair]
         model  = load_model(pair)
         result = model.translate(req.text.strip())
+
+        # Detokenize if needed (Punjabi models)
+        if cfg["detokenize"]:
+            result = detokenize(result)
+
         return {"success": True, "translation": result, "pair": pair}
     except Exception as e:
         print(f"Error [{pair}]: {e}")
-        return {"success": False, "translation": str(e)}
+        return {"success": False, "translation": str(e)}