chatbot / test.py
ogx786's picture
Update test.py
ff31cf6 verified
Raw
History Blame Contribute Delete
974 Bytes
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
local_path = "./local_model"
print("Loading model...")
tokenizer = AutoTokenizer.from_pretrained(local_path)
model = AutoModelForCausalLM.from_pretrained(local_path, torch_dtype=torch.float32)
print("Model loaded.\n")
def transliterate(urdu_text: str) -> str:
prompt = f"""### Instruction:
Transliterate Urdu to Roman Urdu.
### Input:
{urdu_text}
### Response:
"""
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(
**inputs,
max_new_tokens=128,
do_sample=False,
pad_token_id=tokenizer.eos_token_id,
)
generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
return generated.split("### Response:")[-1].strip()
with open("input.txt", "r", encoding="utf-8") as f:
urdu_text = f.read().strip()
print(f"Input: {urdu_text}")
result = transliterate(urdu_text)
print(f"Roman Urdu: {result}")