File size: 4,334 Bytes
a9d8415 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | # ==============================================================================
# 🚀 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()
|