greedy_decode.py

import time

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

def get_model():
    model_name = "meta-llama/Llama-3.2-3B-Instruct"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
    model.eval()
    return model, tokenizer

def get_prompt(tokenizer):
    prompt = (
        "Please fix the following code: \n"
        "```python\n"
        "def hello_world():\n"
        "    print(hello world)\n"
        "```\n"
    )

    message = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": prompt},
    ]

    tokenized_message = tokenizer.apply_chat_template(message, tokenize=True, add_generation_prompt=True, return_tensors="pt")
    return tokenized_message

def greedy_decode(model, input_ids, max_new_tokens=25):
    input_ids = input_ids.to(model.device)
    with torch.no_grad():
        for _ in range(max_new_tokens):
            # forward pass entire sequence (prompt + generated tokens so far)
            # no kv cache
            logits = model(input_ids).logits
            next_token_logits = logits[:, -1, :]
            next_token = torch.argmax(next_token_logits, dim=-1).unsqueeze(0) # shape: (1, 1)
            input_ids = torch.cat([input_ids, next_token], dim=-1)
    return input_ids

if __name__ == "__main__":
    model, tokenizer = get_model()
    prompt = get_prompt(tokenizer)
    max_new_tokens = 25

    start = time.perf_counter()
    generation = greedy_decode(model, prompt, max_new_tokens=max_new_tokens)
    elapsed = time.perf_counter() - start

    print(tokenizer.decode(generation[0], skip_special_tokens=True))
    print(f"Greedy decode time: {elapsed:.2f}s")
    throughput = max_new_tokens / elapsed
    print(f"Output throughput: {throughput:.2f} tokens/s")