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}")