File size: 6,647 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 | # 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
def generate_reference_outputs(total_length, output_file, model_name):
# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logger.info(f"Using device: {device}")
# Load model and tokenizer from HuggingFace
config = AutoConfig.from_pretrained(model_name)
# Qwen only: add rope scaling to the config
# https://huggingface.co/Qwen/Qwen2.5-7B-Instruct#processing-long-texts
if "Qwen" in model_name:
config.rope_scaling = {"factor": 4.0, "original_max_position_embeddings": 32768, "type": "yarn"}
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, config=config, device_map="auto")
model.eval()
# Load the book text
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 = tokenizer.encode(text, add_special_tokens=True)[:total_length]
encoded_tokens_tensor = torch.tensor(encoded_tokens, device=device).unsqueeze(0) # Shape [1, seq_len] on 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
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 using HuggingFace model
outputs = model(chunk_tokens.to(device))
logits = outputs.logits
# Compute top-5 predictions
probs = torch.softmax(logits, 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
chunk_next_tokens_tensor = torch.tensor(
chunk_next_tokens[:actual_chunk_size], device=device
) # Move to same 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))
# Save the data - ensure tensors are concatenated and on CPU
data = {
"top5_tokens": torch.cat(all_top5_tokens, dim=0).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}")
def main():
parser = argparse.ArgumentParser(description="Generate reference outputs using HuggingFace models.")
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, required=True, help="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, model_name=args.model)
if __name__ == "__main__":
main()
|