Update app.py

app.py

@@ -2,9 +2,6 @@ import os
 import torch
 import gradio as gr
 from train import CharTokenizer, Seq2Seq, Encoder, Decoder, TransformerTransliterator
-from huggingface_hub import login
-hf_token = os.getenv('HF_TOKEN')
-login(token=hf_token)
 
 # ----------------------
 # 1️⃣ Load LSTM checkpoint
@@ -25,12 +22,16 @@ DEC_HIDDEN_DIM = 256
 NUM_LAYERS_MODEL = 2
 DROPOUT = 0.3
 
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
 encoder = Encoder(len(src_tokenizer), EMBED_DIM, ENC_HIDDEN_DIM, NUM_LAYERS_MODEL, DROPOUT)
 decoder = Decoder(len(tgt_tokenizer), EMBED_DIM, ENC_HIDDEN_DIM, DEC_HIDDEN_DIM, NUM_LAYERS_MODEL, DROPOUT)
-lstm_model = Seq2Seq(encoder, decoder, device='cpu')
+lstm_model = Seq2Seq(encoder, decoder, device=device).to(device)
 lstm_model.load_state_dict(lstm_ckpt['model_state_dict'])
 lstm_model.eval()
 
+print("✅ LSTM model loaded")
+
 # ----------------------
 # 2️⃣ Load Transformer checkpoint
 # ----------------------
@@ -47,71 +48,109 @@ transformer_model = TransformerTransliterator(
     dim_feedforward=512,
     dropout=0.1,
     max_len=100
-)
+).to(device)
 transformer_model.load_state_dict(transformer_ckpt['model_state_dict'])
 transformer_model.eval()
 
+print("✅ Transformer model loaded")
+
 # ----------------------
-# 3️⃣ Load LLaMA 7B (Hugging Face)
+# 3️⃣ Load lightweight LLM (DistilBERT-based or small model)
 # ----------------------
-# from transformers import LlamaForCausalLM, LlamaTokenizer
-
-# llama_model_name = "meta-llama/Llama-2-7b-hf"  # adjust if using local
-# llama_tokenizer = LlamaTokenizer.from_pretrained(llama_model_name)
-# llama_model = LlamaForCausalLM.from_pretrained(llama_model_name, device_map="auto")
-# llama_model.eval()
-
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
-indic_model_name = "mistralai/Mistral-7B-Instruct-v0.3"
-indic_tokenizer = AutoTokenizer.from_pretrained(indic_model_name)
-indic_model = AutoModelForCausalLM.from_pretrained(indic_model_name)
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+# Use a lightweight T5 model instead of Mistral 7B
+try:
+    llm_model_name = "google/flan-t5-small"  # 60M params, ~240MB
+    llm_tokenizer = AutoTokenizer.from_pretrained(llm_model_name)
+    llm_model = AutoModelForSeq2SeqLM.from_pretrained(llm_model_name).to(device)
+    llm_model.eval()
+    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")
+    has_llm = False
 
 # ----------------------
 # 4️⃣ Transliteration Function
 # ----------------------
+@torch.no_grad()
 def transliterate(word):
     word = word.strip()
-
-    # LSTM prediction
-    lstm_pred = lstm_model.translate(word, src_tokenizer, tgt_tokenizer)
-
-    # Transformer prediction (greedy)
-    transformer_pred = transformer_model.translate(word, src_tokenizer, tgt_tokenizer, decoding="greedy")
-
-    # LLaMA prediction
-    # prompt = f"Transliterate this Hindi Roman word to Devanagari: {word}"
-    # inputs = llama_tokenizer(prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
-    # output_ids = llama_model.generate(**inputs, max_new_tokens=50)
-    # llama_pred = llama_tokenizer.decode(output_ids[0], skip_special_tokens=True).replace(prompt, "").strip()
-
-    inputs = indic_tokenizer(f"transliterate to Devanagari: {word}", return_tensors="pt")
-    output = indic_model.generate(**inputs, max_new_tokens=50)
-    llama_pred = indic_tokenizer.decode(output[0], skip_special_tokens=True)
-
-    return lstm_pred, transformer_pred, llama_pred
+    
+    if not word:
+        return "❌ Empty input", "❌ Empty input", "❌ Empty input"
+    
+    try:
+        # LSTM prediction
+        lstm_pred = lstm_model.translate(word, src_tokenizer, tgt_tokenizer)
+    except Exception as e:
+        lstm_pred = f"Error: {str(e)[:50]}"
+    
+    try:
+        # Transformer prediction (greedy)
+        transformer_pred = transformer_model.translate(
+            word, src_tokenizer, tgt_tokenizer, 
+            device=device, decoding="greedy"
+        )
+    except Exception as e:
+        transformer_pred = f"Error: {str(e)[:50]}"
+    
+    # LLM prediction (lightweight T5)
+    if has_llm:
+        try:
+            prompt = f"Transliterate to Devanagari: {word}"
+            inputs = llm_tokenizer(prompt, return_tensors="pt").to(device)
+            output_ids = llm_model.generate(
+                **inputs, 
+                max_length=20,
+                num_beams=2,
+                early_stopping=True
+            )
+            llm_pred = llm_tokenizer.decode(output_ids[0], skip_special_tokens=True)
+            # Clean up: remove the input prompt if it appears in output
+            llm_pred = llm_pred.replace(prompt, "").strip()
+        except Exception as e:
+            llm_pred = f"Error: {str(e)[:50]}"
+    else:
+        llm_pred = "LLM model not loaded (insufficient memory)"
+    
+    return lstm_pred, transformer_pred, llm_pred
 
 # ----------------------
 # 5️⃣ Gradio Interface
 # ----------------------
-iface = gr.Interface(
+demo = gr.Interface(
     fn=transliterate,
-    inputs=gr.Textbox(label="Input Hindi Roman Word"),
+    inputs=gr.Textbox(
+        label="Input Hindi Roman Word",
+        placeholder="e.g., namaste, dhanyavaad, bharat",
+        lines=1
+    ),
     outputs=[
-        gr.Textbox(label="LSTM Prediction"),
-        gr.Textbox(label="Transformer Prediction"),
-        gr.Textbox(label="Mistral 7B Prediction")
+        gr.Textbox(label="LSTM Prediction", interactive=False),
+        gr.Textbox(label="Transformer Prediction", interactive=False),
+        gr.Textbox(label="Flan-T5 Small Prediction", interactive=False)
     ],
     title="Hindi Roman to Devanagari Transliteration",
-    description="Enter a Hindi Roman word and get predictions from LSTM, Transformer, and Mistral 7B models."
+    description="Compare three models: LSTM, Transformer, and Flan-T5.\nEnter a Hindi Roman word to get transliteration predictions.",
+    examples=[
+        ["namaste"],
+        ["dhanyavaad"],
+        ["bharat"],
+        ["mumbai"],
+        ["hindustan"],
+        ["pranaam"]
+    ],
+    allow_flagging="never"
 )
 
-print('Hello')
-iface.launch(
-    share=True,
-    debug=True,
-    inbrowser=False,
-    server_name="0.0.0.0",
-    server_port=7860,
-    block=True
-)
+if __name__ == "__main__":
+    print("🚀 Starting Gradio interface...")
+    demo.launch(
+        share=False,
+        debug=False,
+        server_name="0.0.0.0",
+        server_port=7860
+    )