Update app.py

app.py

@@ -1,12 +1,24 @@
+import os
 import torch
 import gradio as gr
 from train import CharTokenizer, Seq2Seq, Encoder, Decoder, TransformerTransliterator
 
+# ----------------------
+# ⿡ Load LSTM checkpoint
+# ----------------------
+lstm_ckpt_path = "lstm_transliterator.pt"
+lstm_ckpt = torch.load(lstm_ckpt_path, map_location='cpu')
 
+src_vocab = lstm_ckpt['src_vocab']
+tgt_vocab = lstm_ckpt['tgt_vocab']
 
+src_tokenizer = CharTokenizer(vocab=src_vocab)
+tgt_tokenizer = CharTokenizer(vocab=tgt_vocab)
 
-# ----------------------
-# 1️⃣ Load LSTM checkpoint
+# Reconstruct LSTM model architecture
+EMBED_DIM = 256
+ENC_HIDDEN_DIM = 256
+DEC_HIDDEN_DIM = 256
 NUM_LAYERS_MODEL = 2
 DROPOUT = 0.3
 
@@ -21,8 +33,18 @@ lstm_model.eval()
 print("✅ LSTM model loaded")
 
 # ----------------------
-# 2️⃣ Load Transformer checkpoint
+# ⿢ Load Transformer checkpoint
 # ----------------------
+transformer_ckpt_path = "transformer_transliterator.pt"
+transformer_ckpt = torch.load(transformer_ckpt_path, map_location='cpu')
+
+transformer_model = TransformerTransliterator(
+    src_vocab_size=len(src_tokenizer),
+    tgt_vocab_size=len(tgt_tokenizer),
+    d_model=256,
+    nhead=8,
+    num_encoder_layers=2,
+    num_decoder_layers=2,
     dim_feedforward=512,
     dropout=0.1,
     max_len=100
@@ -33,7 +55,7 @@ transformer_model.eval()
 print("✅ Transformer model loaded")
 
 # ----------------------
-# 3️⃣ Load lightweight LLM (DistilBERT-based or small model)
+# ⿣ Load lightweight LLM (DistilBERT-based or small model)
 # ----------------------
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
@@ -46,12 +68,12 @@ try:
     print("✅ LLM model loaded (Flan-T5 Small)")
     has_llm = True
 except Exception as e:
-    print(f"⚠️ LLM loading failed: {e}")
-    print("⚠️ Will use only LSTM and Transformer models")
+    print(f"⚠ LLM loading failed: {e}")
+    print("⚠ Will use only LSTM and Transformer models")
     has_llm = False
 
 # ----------------------
-# 4️⃣ Transliteration Function
+# ⿤ Transliteration Function
 # ----------------------
 @torch.no_grad()
 def transliterate(word):
@@ -78,7 +100,7 @@ def transliterate(word):
     # LLM prediction (lightweight T5)
     if has_llm:
         try:
-            prompt = f"Transliterate the romanized Hindi word to Devanagari script: {word}"
+            prompt = f"Transliterate the Romanized Hindi word to Devanagari script: {word}"
             inputs = llm_tokenizer(prompt, return_tensors="pt").to(device)
             output_ids = llm_model.generate(
                 **inputs, 
@@ -97,7 +119,7 @@ def transliterate(word):
     return lstm_pred, transformer_pred, llm_pred
 
 # ----------------------
-# 5️⃣ Gradio Interface
+# ⿥ Gradio Interface
 # ----------------------
 demo = gr.Interface(
     fn=transliterate,
@@ -124,11 +146,11 @@ demo = gr.Interface(
     allow_flagging="never"
 )
 
-if __name__ == "__main__":
+if _name_ == "_main_":
     print("🚀 Starting Gradio interface...")
     demo.launch(
         share=False,
         debug=False,
         server_name="0.0.0.0",
         server_port=7860
-    )
+    )