-
Notifications
You must be signed in to change notification settings - Fork 0
/
eval_no_vectors.py
96 lines (64 loc) · 2.82 KB
/
eval_no_vectors.py
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
from evaluator import ArithmeticsEvaluator
from args import TrainingArguments, DataTrainingArguments, ArgumentParser
from arithmetics import PromptArithmeticsConfig
from tasks import Preprocessor, AutoTask
from utils import get_task_prompts
import torch
from trainer import MultiTaskSeq2SeqTrainer
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, DataCollatorForSeq2Seq
from peft import TaskType, PromptTuningConfig, get_peft_model
import os
import numpy as np
from datetime import datetime
import wandb
timestamp = datetime.now().strftime("%m%d%Y%H%M%S")
parser = ArgumentParser(
(TrainingArguments, DataTrainingArguments, PromptArithmeticsConfig)
)
training_args, data_args, pa_config = parser.parse_toml_file("configs/addition.toml")
data_args.dataset_names = sorted(data_args.dataset_names)
# cross origin evaluation on datasets
os.environ["WANDB_PROJECT"] = training_args.wandb_project
tokenizer = AutoTokenizer.from_pretrained(
data_args.data_tokenizer_name_or_path, model_max_length=512, use_fast=True
)
model = AutoModelForSeq2SeqLM.from_pretrained(training_args.model_name_or_path)
peft_config = PromptTuningConfig(
task_type=TaskType.SEQ_2_SEQ_LM,
num_virtual_tokens=pa_config.num_virtual_tokens,
)
model = get_peft_model(model, peft_config)
model.base_model.generation_config.max_new_tokens = data_args.max_target_length
preprocessor = Preprocessor(data_args.dataset_names, data_args, training_args)
_, valid_datasets, test_datasets = preprocessor.get_data()
task_prompts = get_task_prompts(pa_config, data_args.dataset_names)
for origin_prompt in task_prompts:
training_args.origin_prompt_name = origin_prompt
for dataset_name in data_args.dataset_names:
training_args.train_dataset_names = [dataset_name]
mnli_weights = torch.load(f"soft_prompts/{origin_prompt}/mnli.bin")
qnli_weights = torch.load(f"soft_prompts/{origin_prompt}/qnli.bin")
model.prompt_encoder.default.embedding.weight = torch.nn.Parameter(
qnli_weights + mnli_weights
)
data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, return_tensors="pt")
compute_metrics = AutoTask.get(dataset_name).get_compute_metrics(tokenizer)
trainer = MultiTaskSeq2SeqTrainer(
model=model,
tokenizer=tokenizer,
args=training_args,
eval_dataset=valid_datasets,
data_collator=data_collator,
compute_metrics=compute_metrics,
)
print(
trainer.evaluate(
eval_dataset=test_datasets[dataset_name], metric_key_prefix="test"
)
)
save_name = (
f"./saves/prompt_tuning_{timestamp}_{dataset_name}_{origin_prompt}_best"
)
model.save_pretrained(save_name)
if wandb.run is not None:
wandb.finish()