Upload inference.py with huggingface_hub

inference.py

@@ -0,0 +1,128 @@
+# ==============================================================================
+# 🚀 ViuTranslate — Interactive Translation CLI
+# ==============================================================================
+# Run: python inference.py
+# Type any sentence in English or Hindi to get real-time neural translation.
+# ==============================================================================
+
+import os
+import sys
+import torch
+from tokenizers import Tokenizer
+from huggingface_hub import hf_hub_download
+
+# UTF-8 encoding
+if hasattr(sys.stdout, "reconfigure"):
+    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
+
+cur_dir = os.path.dirname(os.path.abspath(__file__)) if "__file__" in locals() else os.getcwd()
+if cur_dir not in sys.path:
+    sys.path.insert(0, cur_dir)
+
+from model import ViuAI
+from config import ViuAIConfig
+
+REPO_ID = "ViuAI/ViuTranslate"
+EOT_ID = 64002
+
+def load_model():
+    print("=" * 75)
+    print(f"🚀 Initializing ViuTranslate-500M from {REPO_ID}...")
+    print("=" * 75)
+    
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"• Hardware Device: {device.type.upper()}")
+    
+    # Checkpoint candidates
+    ckpt_candidates = [
+        "viutranslate_final.pt",
+        "checkpoints/viutranslate_final.pt",
+        os.path.join(cur_dir, "viutranslate_final.pt")
+    ]
+    ckpt_path = None
+    for c in ckpt_candidates:
+        if os.path.exists(c):
+            ckpt_path = c
+            break
+            
+    if ckpt_path is None:
+        print(f"📥 Downloading viutranslate_final.pt from {REPO_ID}...")
+        ckpt_path = hf_hub_download(repo_id=REPO_ID, filename="viutranslate_final.pt")
+        
+    tok_candidates = [
+        "tokenizer.json",
+        os.path.join(cur_dir, "tokenizer.json")
+    ]
+    tok_path = None
+    for tc in tok_candidates:
+        if os.path.exists(tc):
+            tok_path = tc
+            break
+            
+    if tok_path is None:
+        print(f"📥 Downloading tokenizer.json from {REPO_ID}...")
+        tok_path = hf_hub_download(repo_id=REPO_ID, filename="tokenizer.json")
+        
+    tokenizer = Tokenizer.from_file(tok_path)
+    
+    cfg = ViuAIConfig(vocab_size=64003, context_length=2048)
+    model = ViuAI(cfg).to(device)
+    
+    state = torch.load(ckpt_path, map_location=device, weights_only=False)
+    weights = state.get("model_state_dict", state)
+    model.load_state_dict(weights, strict=False)
+    model.eval()
+    
+    print("✅ ViuTranslate Engine loaded and ready for inference!\n")
+    return model, tokenizer, device
+
+@torch.no_grad()
+def translate(model, tokenizer, device, text: str, mode: str = "direct") -> str:
+    text = text.strip()
+    if mode == "direct":
+        prompt = f"<|user|>\n{text}<|endofturn|>\n<|assistant|>\n"
+    elif mode == "to_hi":
+        prompt = f"<|user|>\nTranslate to Hindi: '{text}'<|endofturn|>\n<|assistant|>\n"
+    elif mode == "to_en":
+        prompt = f"<|user|>\nTranslate to English: '{text}'<|endofturn|>\n<|assistant|>\n"
+    else:
+        prompt = f"<|user|>\n{text}<|endofturn|>\n<|assistant|>\n"
+
+    input_ids = torch.tensor([tokenizer.encode(prompt).ids], dtype=torch.long, device=device)
+    prompt_len = input_ids.shape[1]
+    
+    out = model.generate(
+        input_ids,
+        max_new_tokens=150,
+        temperature=0.2,
+        top_p=0.9,
+        repetition_penalty=1.15,
+        eos_token_id=EOT_ID
+    )
+    gen_tokens = out[0][prompt_len:].tolist()
+    if EOT_ID in gen_tokens:
+        gen_tokens = gen_tokens[:gen_tokens.index(EOT_ID)]
+        
+    return tokenizer.decode(gen_tokens).strip()
+
+def interactive_loop():
+    model, tokenizer, device = load_model()
+    print("💡 Enter text to translate (Google Translate style). Type 'exit' or 'quit' to stop.\n")
+    
+    while True:
+        try:
+            inp = input("📝 [Input]: ").strip()
+            if not inp:
+                continue
+            if inp.lower() in ["exit", "quit", "q"]:
+                print("👋 Exiting ViuTranslate.")
+                break
+                
+            out = translate(model, tokenizer, device, inp, mode="direct")
+            print(f"🌐 [ViuTranslate]: {out}\n")
+        except (KeyboardInterrupt, EOFError):
+            print("\n👋 Exiting ViuTranslate.")
+            break
+
+if __name__ == "__main__":
+    interactive_loop()