import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model_name = "Navyaforaa/LitGram-1.5B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16,
device_map="auto"
)
messages = [
{
"role": "system",
"content": "You are LitGram, a specialized model for English literature and English grammar. Give accurate, text-grounded, clearly explained answers."
},
{
"role": "user",
"content": "Scan the metre of Shakespeare's Sonnet 18: 'Shall I compare thee to a summer's day?'"
}
]
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
outputs = model.generate(
**inputs,
max_new_tokens=512,
temperature=0.3,
repetition_penalty=1.15
)
print(tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True))