File size: 8,346 Bytes
babe29f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 | # SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
import argparse
import bz2
import os
import torch
from loguru import logger
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
from models.tt_transformers.tt.model_config import ModelArgs
def generate_reference_outputs(total_length, output_file, hf_model_name=None):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logger.info(f"Using device: {device}")
if hf_model_name:
# HuggingFace path
tokenizer = AutoTokenizer.from_pretrained(hf_model_name)
config = AutoConfig.from_pretrained(hf_model_name)
# Qwen only: add rope scaling to the config
# https://huggingface.co/Qwen/Qwen2.5-7B-Instruct#processing-long-texts
if "Qwen" in hf_model_name:
config.rope_scaling = {"factor": 4.0, "original_max_position_embeddings": 32768, "type": "yarn"}
# Text-only ports of multimodal checkpoints (e.g. EXAONE-4.5): the composite
# config class only maps to an image-text-to-text model.
from transformers import AutoModelForImageTextToText
model_cls = AutoModelForCausalLM
if type(config) not in AutoModelForCausalLM._model_mapping and type(config) in (
AutoModelForImageTextToText._model_mapping
):
model_cls = AutoModelForImageTextToText
model = model_cls.from_pretrained(
hf_model_name, config=config, torch_dtype=torch.float32 if device == "cpu" else None, device_map="auto"
)
model.eval()
model_args = ModelArgs(mesh_device=None, cache_hf=True)
else:
# Original path - load reference model
model_args = ModelArgs(mesh_device=None, cache_hf=True)
model_args.max_seq_len = total_length
tokenizer = model_args.tokenizer
assert tokenizer is not None, "Tokenizer must be provided for non-dummy weights"
reference_model = model_args.reference_transformer(load_checkpoint=True, wrap=False)
reference_model.to(device) # Move model to device
reference_model.eval() # Set to evaluation mode
embd = reference_model.model.embed_tokens
embd.to(device) # Move embedding to device
# Load the book text and encode tokens
current_file_path = os.path.abspath(__file__)
current_file_dir = os.path.dirname(current_file_path)
prompt_file = os.path.join(current_file_dir, "tale-of-two-cities.txt.bz2")
with bz2.open(prompt_file, "rt", encoding="utf-8") as f:
text = f.read()
# Encode text to tokens
encoded_tokens = model_args.encode_prompt(text, instruct=False)
encoded_tokens_tensor = torch.tensor(encoded_tokens, device=device).unsqueeze(0) # Move to device
print(f"{'Progress':<15}{'Correct':<8}{'Actual':<15}{'Top 5 Predictions':<75}")
print("-" * 113)
# Initialize lists to store results
all_top1_correct = []
all_top5_correct = []
all_top5_tokens = []
segment_accuracies = []
chunk_size = 1024
with torch.no_grad():
for chunk_start in range(0, total_length - 1, chunk_size):
chunk_end = min(chunk_start + chunk_size, total_length)
# Get input and target chunks, ensuring they have matching lengths
chunk_tokens = encoded_tokens_tensor[:, chunk_start:chunk_end]
chunk_next_tokens = encoded_tokens[chunk_start + 1 : chunk_end + 1]
actual_chunk_size = min(len(chunk_tokens[0]), len(chunk_next_tokens))
# Trim input chunk if needed
chunk_tokens = chunk_tokens[:, :actual_chunk_size]
# Process chunk based on model type
chunk_tokens = chunk_tokens.to(device)
if hf_model_name:
outputs = model(chunk_tokens)
ref_output = outputs.logits
else:
pt_decode_input = embd(chunk_tokens).view(1, actual_chunk_size, -1)
ref_output = reference_model(pt_decode_input, start_pos=chunk_start)
# Compute top-5 predictions
probs = torch.softmax(ref_output, dim=-1)
_, chunk_top5_tokens = torch.topk(probs, k=5, dim=-1) # Shape: [1, chunk_size, 5]
chunk_top5_tokens = chunk_top5_tokens.squeeze(0) # Shape: [chunk_size, 5]
# Get next tokens tensor, ensuring same length as predictions
chunk_next_tokens_tensor = torch.tensor(chunk_next_tokens[:actual_chunk_size], device=device)
# Calculate correctness
chunk_top1_correct = chunk_top5_tokens[:, 0] == chunk_next_tokens_tensor
chunk_top5_correct = torch.any(chunk_top5_tokens == chunk_next_tokens_tensor.unsqueeze(1), dim=1)
# Store results
all_top1_correct.extend(chunk_top1_correct.tolist())
all_top5_correct.extend(chunk_top5_correct.tolist())
all_top5_tokens.append(chunk_top5_tokens)
# Print predictions for this chunk
for i in range(len(chunk_next_tokens)):
global_pos = chunk_start + i
next_token = chunk_next_tokens[i]
sanitize = lambda x: x.replace("\n", "").replace("\r", "").replace("\x0c", "")
actual_token = sanitize(tokenizer.decode([next_token]))
top5_tokens = [sanitize(tokenizer.decode([t.item()])) for t in chunk_top5_tokens[i]]
correct = "x" if chunk_top1_correct[i] else ("-" if chunk_top5_correct[i] else " ")
top5_str = " ".join(f"{t:<14}" for t in top5_tokens)
progress_str = f"{global_pos+1}/{total_length-1}"
print(f"{progress_str:<15}{correct:<8}{actual_token:<15}{top5_str}")
# Calculate and store segment accuracies every 100 tokens
if (global_pos + 1) % 100 == 0 or global_pos == total_length - 2:
start_idx = (global_pos // 100) * 100
end_idx = min(start_idx + 100, len(all_top1_correct))
segment_top1_acc = sum(all_top1_correct[start_idx:end_idx]) / (end_idx - start_idx) * 100
segment_top5_acc = sum(all_top5_correct[start_idx:end_idx]) / (end_idx - start_idx) * 100
if len(segment_accuracies) <= global_pos // 100:
segment_accuracies.append((segment_top1_acc, segment_top5_acc))
# Concatenate all top5 tokens into a single tensor
all_top5_tokens = torch.cat(all_top5_tokens, dim=0) # Shape: [total_tokens, 5]
# Move tensors back to CPU before saving
data = {
"top5_tokens": all_top5_tokens.cpu(),
"reference_tokens": encoded_tokens_tensor[:, :total_length].clone().cpu(),
}
torch.save(data, output_file)
logger.info(f"Saved reference outputs to {output_file}")
# Print all segment accuracy summaries as a table
print("\nSegment Accuracy Summaries:")
print(f"{'Tokens':<15}{'Top-1 Accuracy':<20}{'Top-5 Accuracy':<20}")
print("-" * 55)
for i, (top1_acc, top5_acc) in enumerate(segment_accuracies):
start_token = i * 100 + 1
end_token = min((i + 1) * 100, total_length)
print(f"{f'{start_token}-{end_token}':<15}{f'{top1_acc:.2f}%':<20}{f'{top5_acc:.2f}%':<20}")
# Calculate overall accuracy
overall_top1_acc = sum(acc[0] for acc in segment_accuracies) / len(segment_accuracies)
overall_top5_acc = sum(acc[1] for acc in segment_accuracies) / len(segment_accuracies)
print("-" * 55)
print(f"{'Overall':<15}{f'{overall_top1_acc:.2f}%':<20}{f'{overall_top5_acc:.2f}%':<20}")
# New main function with argparse
def main():
parser = argparse.ArgumentParser(description="Generate reference outputs for LLaMA accuracy testing.")
parser.add_argument("--total_length", type=int, default=1024, help="Total length of tokens to process")
parser.add_argument(
"--output_file", type=str, default="reference_outputs.pt", help="Output file path for reference data"
)
parser.add_argument(
"--model", type=str, help="Optional: HuggingFace model name (e.g., 'meta-llama/Llama-3.1-8B-Instruct')"
)
args = parser.parse_args()
generate_reference_outputs(total_length=args.total_length, output_file=args.output_file, hf_model_name=args.model)
if __name__ == "__main__":
main()
|