garystafford/Llama-3.2-3B-Instruct-lora-nvidia-blackwell

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

device = "cuda:0" if torch.cuda.is_available() else "cpu"

HF_TOKEN = "<YOUR_HF_TOKEN_GOES_HERE>"
base_model_id = "meta-llama/Llama-3.2-3B-Instruct"
peft_model_id = "garystafford/Llama-3.2-3B-Instruct-lora-nvidia-blackwell"

base_model = AutoModelForCausalLM.from_pretrained(base_model_id, token=HF_TOKEN)
model = PeftModel.from_pretrained(base_model, peft_model_id).to(device)
tokenizer = AutoTokenizer.from_pretrained(base_model_id, token=HF_TOKEN)

test_prompt = [
  {
      "role": "user",
      "content": "Describe the NVIDIA Blackwell architecture.",
  }
]

inputs = tokenizer.apply_chat_template(
    test_prompt,
    tokenize=True,
    add_generation_prompt=True,
    return_tensors="pt",
).to(device)

output = model.generate(
    input_ids=inputs,
    max_new_tokens=128,
    temperature=0.1,
    pad_token_id=tokenizer.pad_token_id,
).to(device)

output = tokenizer.decode(output[0], skip_special_tokens=True)
output = output.split('assistant\n\n')[1].strip()
print(output)

@misc{vonwerra2022trl,
	title        = {{TRL: Transformer Reinforcement Learning}},
	author       = {Leandro von Werra and Younes Belkada and Lewis Tunstall and Edward Beeching and Tristan Thrush and Nathan Lambert and Shengyi Huang and Kashif Rasul and Quentin Gallou{\'e}dec},
	year         = 2020,
	journal      = {GitHub repository},
	publisher    = {GitHub},
	howpublished = {\url{https://github.com/huggingface/trl}}
}