import torch
from transformers import AutoTokenizer, TextStreamer, AutoModelForCausalLM

model_path = "cognitivecomputations/Quiet-STaR-Base"
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    device_map="auto",
    low_cpu_mem_usage=True,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
)

tokenizer = AutoTokenizer.from_pretrained(model_path)
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

prompt = "Hello my name is"

tokens = tokenizer(
    prompt, 
    return_tensors='pt'
).input_ids.cuda()

generation_output = model.generate(
    tokens, 
    streamer=streamer,
    max_new_tokens=512,
)