File size: 2,252 Bytes
24ebd71 | 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 | from __future__ import annotations
import json
import time
import torch
from snip_common import (
ARTIFACT_DIR,
DATA_DIR,
parameter_count,
perplexity,
read_texts,
texts_to_blocks,
)
from transformers import GPT2LMHeadModel, PreTrainedTokenizerFast
PROMPTS = [
"Once upon a time",
"The little robot discovered",
"Jacob opened the castle door and",
]
def main() -> None:
tokenizer = PreTrainedTokenizerFast.from_pretrained(ARTIFACT_DIR)
model = GPT2LMHeadModel.from_pretrained(ARTIFACT_DIR)
model.eval()
dataset = texts_to_blocks(read_texts(DATA_DIR / "eval.jsonl"), tokenizer)
losses: list[float] = []
started = time.perf_counter()
with torch.no_grad():
for index in range(min(len(dataset), 200)):
row = dataset[index]
input_ids = torch.tensor([row["input_ids"]], dtype=torch.long)
outputs = model(input_ids=input_ids, labels=input_ids)
losses.append(float(outputs.loss))
elapsed = time.perf_counter() - started
mean_loss = sum(losses) / len(losses)
samples = []
for prompt in PROMPTS:
encoded = tokenizer(prompt, return_tensors="pt")
with torch.no_grad():
output = model.generate(
**encoded,
max_new_tokens=64,
do_sample=True,
temperature=0.85,
top_k=40,
top_p=0.92,
repetition_penalty=1.08,
pad_token_id=tokenizer.pad_token_id,
eos_token_id=tokenizer.eos_token_id,
)
samples.append(
{
"prompt": prompt,
"completion": tokenizer.decode(output[0], skip_special_tokens=True),
}
)
results = {
"model": "SNIP-0.4M",
"parameters": parameter_count(model),
"eval_blocks": len(losses),
"eval_loss": mean_loss,
"perplexity": perplexity(mean_loss),
"eval_seconds": elapsed,
"samples": samples,
}
(ARTIFACT_DIR / "evaluation.json").write_text(
json.dumps(results, indent=2),
encoding="utf-8",
)
print(json.dumps(results, indent=2))
if __name__ == "__main__":
main()
|