import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# 1. Load the model and its tokenizer
model_id = "HuggingFaceTB/SmolLM2-135M-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)

# Move to GPU if available, else keep on CPU
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

# 2. Format your text input using a Chat Template
messages = [
    {"role": "user", "content": "Explain what a 4D Discrete Cosine Transform does in one sentence."}
]

# --- Your previous tokenization step ---
inputs = tokenizer.apply_chat_template(
    messages, 
    add_generation_prompt=True, 
    return_tensors="pt"
).to(device)

# --- THE FIX ---
# Check if inputs is a dictionary/BatchEncoding; grab the underlying tensor
if isinstance(inputs, dict) or hasattr(inputs, "data"):
    input_ids = inputs["input_ids"]
else:
    input_ids = inputs

# Pass the tensor directly to generate
outputs = model.generate(
    input_ids,  # Use input_ids tensor instead of the dictionary wrapper
    max_new_tokens=60, 
    temperature=0.2, 
    do_sample=True
)

# Use input_ids.shape[-1] to strip away the prompt when decoding
response = tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True)

print("Model Output:", response)