nanochat-students/rl-d20

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


model_name = "nanochat-students/rl-d20"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True).to(device)
model.eval()

conversation = [
    {"role": "user", "content": "Hello, who are you?"},
]
rendered = tokenizer.apply_chat_template(
    conversation,
    tokenize=False,
    add_generation_prompt=True,
)
model_inputs = tokenizer([rendered], return_tensors="pt").to(model.device)

generated = model.generate(**model_inputs, max_new_tokens=256)
output_ids = generated[0, model_inputs.input_ids.shape[1]:]
print(tokenizer.decode(output_ids, skip_special_tokens=True))