Views
No views yet
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
repo_id = "adhitia17/idmt"
print(f"Loading tokenizer and model from {repo_id}...")
tokenizer = AutoTokenizer.from_pretrained(repo_id)
model = AutoModelForSeq2SeqLM.from_pretrained(repo_id)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
print(f"Model loaded to device: {device}")
def generate_response(input_text, task_prefix):
"""Generates a response from the model for a given task."""
full_input = f"{task_prefix}: {input_text}"
print(f"\nInput ({task_prefix}): {full_input}")
input_ids = tokenizer(full_input, return_tensors="pt").input_ids.to(device)
outputs = model.generate(
input_ids,
max_length=256,
num_beams=5,
early_stopping=True
)
decoded_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(f"Output: {decoded_output}")
return decoded_output
print("\nInference examples complete.")