import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load tokenizer and model
model_name = "baban/MT_En_Hindi" 
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto")

source_text = "The weather is nice today."
prompt = f"Translate the following English sentence to Hindi:\n{source_text}"

messages = [
    {"role": "user", "content": prompt}
]

# Tokenize the formatted input
input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device)


with torch.no_grad():
    output_ids = model.generate(
      input_ids=input_ids,
      max_new_tokens=100,
      do_sample=False
  )


# Decode and print only the new tokens (the response)
response = tokenizer.decode(output_ids[0][input_ids.shape[-1]:], skip_special_tokens=True)
print("\n=== Translation ===")
print(response)