4_load_directly_chat.py

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from models import MODEL_IDS

model_id = MODEL_IDS["llama_3.2_1b_instruct"]

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)

# Ensure the tokenizer has a pad token
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token  # <|end_of_text|>

# Move model to GPU if available
device = "mps" if torch.backends.mps.is_available(
) else "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)

# Construct the Instruct Model Prompt
# Visit: https://www.llama.com/docs/model-cards-and-prompt-formats/llama3_1/#-instruct-model-prompt-
user_input = "What is the capital of France?"
prompt = f"""
<|begin_of_text|><|start_header_id|>system<|end_header_id|>

You are a helpful assistant<|eot_id|><|start_header_id|>user<|end_header_id|>

{user_input}<|eot_id|><|start_header_id|>assistant<|end_header_id|>
"""

# Tokenize input with attention mask
inputs = tokenizer(prompt, return_tensors="pt",
                   padding=True)  # input_ids, attention_mask
input_ids = inputs.input_ids.to(device)
attention_mask = inputs.attention_mask.to(device)

# Generate output with attention mask
output = model.generate(
    input_ids, attention_mask=attention_mask, max_length=128, do_sample=True)

# Decode the response with special tokens included
response = tokenizer.decode(output[0], skip_special_tokens=False)

# Define markers for assistant's response
start_marker = "<|start_header_id|>assistant<|end_header_id|>"
end_marker = "<|eot_id|>"

# Extract the assistant's reply
if start_marker in response:
    response = response.split(start_marker, 1)[-1].strip()
if end_marker in response:
    response = response.split(end_marker, 1)[0].strip()

# Print the cleaned response
print("\n" + response)