| 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.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.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() |
