File size: 5,450 Bytes
c1b8df7 | 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 | """
PACE Actor Model – Test & Evaluation
Loads the fine-tuned checkpoint, evaluates accuracy / perplexity on a held-out
test split, and runs a code-generation demo with a sample prompt.
"""
import math
import torch
from torch.utils.data import DataLoader, random_split
from transformers import AutoTokenizer, AutoModelForCausalLM
from tqdm import tqdm
from dataset import CoderActorDataset
def evaluate(model, dataloader, device):
"""
Runs the model over every batch in *dataloader* and returns:
avg_loss – mean cross-entropy loss across all batches
perplexity – exp(avg_loss)
accuracy – token-level prediction accuracy (ignoring padding / -100 labels)
"""
model.eval()
total_loss = 0.0
total_correct = 0
total_tokens = 0
with torch.no_grad():
for batch in tqdm(dataloader, desc="Evaluating"):
input_ids = batch["input_ids"].to(device)
labels = batch["labels"].to(device)
attention_mask = batch["attention_mask"].to(device)
outputs = model(
input_ids=input_ids,
attention_mask=attention_mask,
labels=labels,
)
total_loss += outputs.loss.item()
# ── Token-level accuracy ──────────────────────────────
# Causal LM: logits are shifted so that logits[t] predicts token[t+1].
# HuggingFace already does this shift internally for the loss,
# but we need to do it ourselves for accuracy.
shift_logits = outputs.logits[:, :-1, :] # (B, T-1, V)
shift_labels = labels[:, 1:] # (B, T-1)
predictions = shift_logits.argmax(dim=-1) # (B, T-1)
# Only count positions where the label is not -100 (padding)
valid_mask = shift_labels != -100
total_correct += (
(predictions[valid_mask] == shift_labels[valid_mask]).sum().item()
)
total_tokens += valid_mask.sum().item()
avg_loss = total_loss / len(dataloader)
perplexity = math.exp(avg_loss)
accuracy = (total_correct / total_tokens * 100) if total_tokens > 0 else 0.0
return avg_loss, perplexity, accuracy
def generate_code(model, tokenizer, prompt, device, max_new_tokens=200):
"""
Generates Python code from a natural-language or code *prompt*.
Returns the decoded string (prompt + completion).
"""
model.eval()
inputs = tokenizer(prompt, return_tensors="pt").to(device)
with torch.no_grad():
output_ids = model.generate(
**inputs,
max_new_tokens=max_new_tokens,
do_sample=True,
temperature=0.7,
top_p=0.95,
repetition_penalty=1.15,
pad_token_id=tokenizer.eos_token_id,
)
return tokenizer.decode(output_ids[0], skip_special_tokens=True)
def main():
# ==========================================
# 1. HARDWARE SETUP
# ==========================================
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}\n")
# ==========================================
# 2. LOAD FINE-TUNED MODEL & TOKENIZER
# ==========================================
model_dir = "masteries/coding/models/actor_v1"
print(f"Loading fine-tuned model from: {model_dir}")
tokenizer = AutoTokenizer.from_pretrained(model_dir)
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(model_dir).to(device)
print(
f"Model loaded – {sum(p.numel() for p in model.parameters()):,} parameters\n"
)
# ==========================================
# 3. PREPARE TEST SPLIT
# ==========================================
full_dataset = CoderActorDataset(
parquet_path="masteries/coding/data/raw/stack_python_157k.parquet",
tokenizer=tokenizer,
max_length=512,
)
# Use 1 % of the data as a test set (fast but representative)
test_size = max(1, int(len(full_dataset) * 0.01))
_, test_dataset = random_split(
full_dataset,
[len(full_dataset) - test_size, test_size],
generator=torch.Generator().manual_seed(42),
)
test_loader = DataLoader(test_dataset, batch_size=8, shuffle=False)
print(f"Test split: {len(test_dataset)} samples ({len(test_loader)} batches)\n")
# ==========================================
# 4. EVALUATION – Loss / Perplexity / Accuracy
# ==========================================
print("=" * 60)
print(" EVALUATION RESULTS")
print("=" * 60)
avg_loss, perplexity, accuracy = evaluate(model, test_loader, device)
print(f"\n Average Loss : {avg_loss:.4f}")
print(f" Perplexity : {perplexity:.2f}")
print(f" Token Accuracy: {accuracy:.2f}%\n")
# ==========================================
# 5. CODE GENERATION – Test Prompt
# ==========================================
test_prompt = "def fibonacci(n):"
print("=" * 60)
print(" CODE GENERATION DEMO")
print("=" * 60)
print(f"\n Prompt: {test_prompt}\n")
print("-" * 60)
generated = generate_code(model, tokenizer, test_prompt, device)
print(generated)
print("-" * 60)
print("\nTest complete.")
if __name__ == "__main__":
main()
|