Update app.py

app.py

@@ -185,11 +185,16 @@ def load_model_and_vocabs():
     with open('latin_stoi.json', 'r', encoding='utf-8') as f:
         latin_stoi = json.load(f)
     with open('latin_itos.json', 'r', encoding='utf-8') as f:
-        latin_itos = json.load(f)
+        latin_itos_raw = json.load(f)
+        # Convert string keys to integers
+        latin_itos = {int(k): v for k, v in latin_itos_raw.items()}
+    
     with open('arabic_stoi.json', 'r', encoding='utf-8') as f:
         arabic_stoi = json.load(f)
     with open('arabic_itos.json', 'r', encoding='utf-8') as f:
-        arabic_itos = json.load(f)
+        arabic_itos_raw = json.load(f)
+        # Convert string keys to integers
+        arabic_itos = {int(k): v for k, v in arabic_itos_raw.items()}
     
     # Initialize model
     model = CTCTransliterator(
@@ -198,7 +203,7 @@ def load_model_and_vocabs():
         len(arabic_stoi),
         num_layers=3,
         dropout=0.3,
-        upsample_factor=2  # ← ADD THIS
+        upsample_factor=2
     ).to(device)
     
     # Load trained weights
@@ -216,22 +221,16 @@ def encode_text(text, vocab):
     """Encode text using vocabulary"""
     return torch.tensor([vocab.get(ch, 0) for ch in text.strip()], dtype=torch.long)
 
-def greedy_decode(log_probs, blank_id):
+def greedy_decode(log_probs, blank_id, itos, stoi):
     """
     Decode CTC outputs using greedy decoding.
-
-    Args:
-        log_probs: (T, B, C) - log probabilities from model
-        input_lengths: (B,) - actual lengths of each sequence (optional)
     """
-    # log_probs: (T, B, C)
-    eos_id = arabic_stoi.get('<eos>',len(arabic_stoi)-2)
+    eos_id = stoi.get('<eos>', len(stoi)-2)
     preds = log_probs.argmax(2).T.cpu().numpy()  # (B, T)
     results = []
     raw_results = []
 
     for i, pred in enumerate(preds):
-
         prev = None
         decoded = []
         raw_result = []
@@ -241,9 +240,9 @@ def greedy_decode(log_probs, blank_id):
                 break
             # CTC collapse: skip blanks and repeated characters
             if p != blank_id and p != prev:
-                decoded.append(arabic_itos[p])
+                decoded.append(itos[str(p)])  # Convert to string if needed
             prev = p
-            raw_result.append(arabic_itos[p])
+            raw_result.append(itos[str(p)])
 
         results.append("".join(decoded))
         raw_results.append("".join(raw_result))
@@ -269,7 +268,7 @@ def transliterate_latin_to_arabic(text):
             out = model(src)
         
         # Decode output
-        decoded = greedy_decode(out, blank_id)
+        decoded = greedy_decode(out, blank_id, arabic_itos, arabic_stoi)
         result = decoded[0] if decoded else ""
         
         # Cache the result in Firebase