File size: 4,358 Bytes
b24b632 | 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 | import json
import os
import torch
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm
from transformers import AutoConfig, AutoModelForCausalLM
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
torch.set_default_dtype(torch.float32)
print(f"Using device: {DEVICE}")
PMNET_MODEL_ID = "phasorkinetics/pmnet"
PMNET_CONFIG_PATH = "config_pmnet.json"
PMNET_CKPT = "copy_model_pmnet.safetensors"
PMNET_NO_MEM_CONFIG_PATH = "config_pmnet_no_mem.json"
PMNET_NO_MEM_CKPT = "copy_model_pmnet_no_mem.safetensors"
VOCAB_SIZE = 250
DELIM_TOKEN = 251
SEQ_LENGTHS = [10, 64, 128, 256, 512, 768, 896, 1024, 1088, 1152, 1280]
NUM_SAMPLES = 100
BATCH_SIZE = 10
def load_eval_model(config_path, ckpt_path):
print(f"Loading model from {ckpt_path}...")
with open(config_path, "r", encoding="utf-8") as f:
config_dict = json.load(f)
config = AutoConfig.from_pretrained(PMNET_MODEL_ID, trust_remote_code=True)
config.update(config_dict)
model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)
from safetensors.torch import load_file
state_dict = load_file(ckpt_path)
model.load_state_dict(state_dict, strict=False)
model.float().to(DEVICE)
model.eval()
return model
def evaluate_copy_paste(model, seq_lengths, num_samples, batch_size):
accuracies = []
with torch.no_grad():
for seq_len in seq_lengths:
print(f"Evaluating Sequence Length: {seq_len}...")
total_correct = 0
total_tokens = 0
for _ in range(0, num_samples, batch_size):
current_batch_size = min(batch_size, num_samples - _)
orig_seq = torch.randint(1, VOCAB_SIZE + 1, (current_batch_size, seq_len), device=DEVICE)
delim = torch.full((current_batch_size, 1), DELIM_TOKEN, device=DEVICE)
input_ids = torch.cat([orig_seq, delim, orig_seq], dim=1)
outputs = model(input_ids)
logits = outputs.logits
shifted_logits = logits[:, :-1, :].contiguous()
shifted_labels = input_ids[:, 1:].contiguous()
target_logits = shifted_logits[:, seq_len:, :]
target_labels = shifted_labels[:, seq_len:]
preds = target_logits.argmax(dim=-1)
correct = (preds == target_labels).sum().item()
total_correct += correct
total_tokens += (current_batch_size * seq_len)
acc = (total_correct / total_tokens) * 100.0
accuracies.append(acc)
print(f" -> Accuracy for len {seq_len}: {acc:.2f}%")
return accuracies
def main():
model_pmnet = load_eval_model(PMNET_CONFIG_PATH, PMNET_CKPT)
acc_pmnet = evaluate_copy_paste(model_pmnet, SEQ_LENGTHS, NUM_SAMPLES, BATCH_SIZE)
del model_pmnet
torch.cuda.empty_cache()
model_no_mem = load_eval_model(PMNET_NO_MEM_CONFIG_PATH, PMNET_NO_MEM_CKPT)
acc_no_mem = evaluate_copy_paste(model_no_mem, SEQ_LENGTHS, NUM_SAMPLES, BATCH_SIZE)
del model_no_mem
torch.cuda.empty_cache()
plt.figure(figsize=(8, 6))
plt.plot(SEQ_LENGTHS, acc_pmnet, marker='o', linestyle='-', color='seagreen', linewidth=2.5, label='PMNet (Full Memory)')
plt.plot(SEQ_LENGTHS, acc_no_mem, marker='s', linestyle='--', color='mediumvioletred', linewidth=2.5, label='PMNet (No Memory)')
# plt.axvline(x=128, color='gray', linestyle=':', linewidth=2, label='SWA Window Size (128)')
plt.title('Copy-Paste Task: Accuracy', fontsize=16, fontweight='bold')
plt.xlabel('Sequence Length to Memorize (Bytes)', fontsize=14)
plt.ylabel('Exact Token Accuracy (%)', fontsize=14)
# plt.xscale('log', base=2)
plt.xticks(SEQ_LENGTHS, [str(x) for x in SEQ_LENGTHS], rotation=35, ha='right', fontsize=12)
plt.yticks(np.arange(0, 101, 10), fontsize=12)
plt.ylim(-5, 105)
plt.grid(True, which="both", ls="--", alpha=0.5)
plt.legend(fontsize=12, loc='lower left')
plt.tight_layout()
save_path = "copy_paste.pdf"
plt.savefig(save_path, dpi=300)
print(f"\nEvaluation complete! Graph saved to {save_path}")
if __name__ == "__main__":
main() |