File size: 3,259 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 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | from __future__ import annotations
import json
import torch
from lora_data import encode_examples, split_examples
from peft import PeftModel
from snip_common import ARTIFACT_DIR, perplexity
from train_lora import ADAPTER_DIR, MERGED_DIR
from transformers import AutoModelForCausalLM, PreTrainedTokenizerFast
def score(model, dataset, limit: int = 90) -> float:
model.eval()
losses = []
with torch.no_grad():
for index in range(min(limit, len(dataset))):
row = dataset[index]
input_ids = torch.tensor([row["input_ids"]], dtype=torch.long)
attention_mask = torch.tensor([row["attention_mask"]], dtype=torch.long)
labels = torch.tensor([row["labels"]], dtype=torch.long)
losses.append(
float(
model(
input_ids=input_ids,
attention_mask=attention_mask,
labels=labels,
).loss
)
)
return sum(losses) / len(losses)
def generate(model, tokenizer, prompt: str) -> str:
prefix = f"<bos>User: {prompt}\nAssistant:"
tokenizer.truncation_side = "left"
encoded = tokenizer(
prefix,
return_tensors="pt",
add_special_tokens=False,
truncation=True,
max_length=80,
)
with torch.no_grad():
output = model.generate(
**encoded,
max_new_tokens=48,
do_sample=True,
temperature=0.75,
top_k=35,
top_p=0.9,
repetition_penalty=1.08,
pad_token_id=tokenizer.pad_token_id,
eos_token_id=tokenizer.eos_token_id,
)
return tokenizer.decode(output[0], skip_special_tokens=True)
def main() -> None:
tokenizer = PreTrainedTokenizerFast.from_pretrained(ARTIFACT_DIR)
_, eval_examples = split_examples()
dataset = encode_examples(eval_examples, tokenizer)
base = AutoModelForCausalLM.from_pretrained(ARTIFACT_DIR)
adapter = PeftModel.from_pretrained(
AutoModelForCausalLM.from_pretrained(ARTIFACT_DIR),
ADAPTER_DIR,
)
merged = AutoModelForCausalLM.from_pretrained(MERGED_DIR)
base_loss = score(base, dataset)
adapter_loss = score(adapter, dataset)
merged_loss = score(merged, dataset)
prompt = "Write a tiny story: a careful robot finds a key in a moonlit castle."
results = {
"evaluation_examples": min(90, len(dataset)),
"base_response_loss": base_loss,
"adapter_response_loss": adapter_loss,
"merged_response_loss": merged_loss,
"base_response_perplexity": perplexity(base_loss),
"adapter_response_perplexity": perplexity(adapter_loss),
"improvement_percent": 100 * (base_loss - adapter_loss) / base_loss,
"sample_prompt": prompt,
"base_sample": generate(base, tokenizer, prompt),
"adapter_sample": generate(adapter, tokenizer, prompt),
"merged_sample": generate(merged, tokenizer, prompt),
}
(ADAPTER_DIR / "variant_evaluation.json").write_text(
json.dumps(results, indent=2),
encoding="utf-8",
)
print(json.dumps(results, indent=2))
if __name__ == "__main__":
main()
|