def train_pet_ensemble(model_config: WrapperConfig,
train_config: TrainConfig,
eval_config: EvalConfig,
pattern_ids: List[int],
output_dir: str,
ipet_data_dir: str = None,
repetitions: int = 3,
train_data: List[InputExample] = None,
unlabeled_data: List[InputExample] = None,
eval_data: List[InputExample] = None,
do_train: bool = True,
do_eval: bool = True,
save_unlabeled_logits: bool = False,
seed: int = 42):
"""
Train and evaluate an ensemble of PET models without knowledge distillation.
:param model_config: the model configuration to use
:param train_config: the training configuration to use
:param eval_config: the evaluation configuration to use
:param pattern_ids: the ids of all PVPs to use
:param output_dir: the output directory
:param ipet_data_dir: optional directory containing additional training data for iPET
:param repetitions: the number of training repetitions
:param train_data: the training examples to use
:param unlabeled_data: the unlabeled examples to use
:param eval_data: the evaluation examples to use
:param do_train: whether to perform training
:param do_eval: whether to perform evaluation
:param save_unlabeled_logits: whether logits for unlabeled examples should be saved in a file ``logits.txt``. This
is required for both iPET and knowledge distillation.
:param seed: the random seed to use
"""
results = defaultdict(lambda: defaultdict(list))
set_seed(seed)
for pattern_id in pattern_ids:
for iteration in range(repetitions):
model_config.pattern_id = pattern_id
results_dict = {}
pattern_iter_output_dir = "{}/p{}-i{}".format(
output_dir, pattern_id, iteration)
if os.path.exists(pattern_iter_output_dir):
logger.warning(
f"Path {pattern_iter_output_dir} already exists, skipping it..."
)
continue
if not os.path.exists(pattern_iter_output_dir):
os.makedirs(pattern_iter_output_dir)
wrapper = init_model(model_config)
# Training
if do_train:
if ipet_data_dir:
p = os.path.join(
ipet_data_dir,
'p{}-i{}-train.bin'.format(pattern_id, iteration))
ipet_train_data = InputExample.load_examples(p)
for example in ipet_train_data:
example.logits = None
else:
ipet_train_data = None
results_dict.update(
train_single_model(wrapper,
train_data,
train_config,
eval_config,
ipet_train_data=ipet_train_data,
unlabeled_data=unlabeled_data))
with open(os.path.join(pattern_iter_output_dir, 'results.txt'),
'w') as fh:
fh.write(str(results_dict))
logger.info("Saving trained model at {}...".format(
pattern_iter_output_dir))
wrapper.save(pattern_iter_output_dir)
train_config.save(
os.path.join(pattern_iter_output_dir, 'train_config.json'))
eval_config.save(
os.path.join(pattern_iter_output_dir, 'eval_config.json'))
logger.info("Saving complete")
if save_unlabeled_logits:
logits = evaluate(wrapper, unlabeled_data,
eval_config)['logits']
save_logits(
os.path.join(pattern_iter_output_dir, 'logits.txt'),
logits)
if not do_eval:
wrapper.model = None
wrapper = None
torch.cuda.empty_cache()
# Evaluation
if do_eval:
logger.info("Starting evaluation...")
if not wrapper:
wrapper = TransformerModelWrapper.from_pretrained(
pattern_iter_output_dir)
eval_result = evaluate(wrapper,
eval_data,
eval_config,
priming_data=train_data)
save_predictions(
os.path.join(pattern_iter_output_dir, 'predictions.jsonl'),
wrapper, eval_result)
save_logits(
os.path.join(pattern_iter_output_dir, 'eval_logits.txt'),
eval_result['logits'])
scores = eval_result['scores']
logger.info(
"--- RESULT (pattern_id={}, iteration={}) ---".format(
pattern_id, iteration))
logger.info(scores)
results_dict['test_set_after_training'] = scores
with open(
os.path.join(pattern_iter_output_dir, 'results.json'),
'w') as fh:
json.dump(results_dict, fh)
for metric, value in scores.items():
results[metric][pattern_id].append(value)
wrapper.model = None
wrapper = None
torch.cuda.empty_cache()
if do_eval:
logger.info("=== OVERALL RESULTS ===")
_write_results(os.path.join(output_dir, 'result_test.txt'), results)
else:
logger.info("=== ENSEMBLE TRAINING COMPLETE ===")
def train_pet_ensemble(
model_config: WrapperConfig,
train_config: TrainConfig,
eval_config: EvalConfig,
pattern_ids: List[Union[str, int]],
output_dir: str,
ipet_data_dir: str = None,
repetitions: int = 3,
train_data: List[InputExample] = None,
unlabeled_data: List[InputExample] = None,
dev_data: List[InputExample] = None,
test_data: List[InputExample] = None,
do_train: bool = True,
do_eval: bool = True,
save_unlabeled_logits: bool = False,
seed: int = 42,
overwrite_dir: bool = False,
save_model=False,
local_rank=-1,
):
"""
Train and evaluate an ensemble of PET models without knowledge distillation.
:param model_config: the model configuration to use
:param train_config: the training configuration to use
:param eval_config: the evaluation configuration to use
:param pattern_ids: the ids of all PVPs to use
:param output_dir: the output directory
:param ipet_data_dir: optional directory containing additional training data for iPET
:param repetitions: the number of training repetitions
:param train_data: the training examples to use
:param unlabeled_data: the unlabeled examples to use
:param dev_data: the evaluation examples to use
:param do_train: whether to perform training
:param do_eval: whether to perform evaluation
:param save_unlabeled_logits: whether logits for unlabeled examples should be saved in a file ``logits.txt``. This
is required for both iPET and knowledge distillation.
:param seed: the random seed to use
"""
results = defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
set_seed(seed)
for pattern_id in pattern_ids:
for iteration in range(repetitions):
model_config.pattern_id = pattern_id
results_dict = {}
shots = 0 if train_data is None else len(train_data)
pattern_iter_output_dir = "{}/{}shots-{}-i{}-seed{}".format(
output_dir, shots, pattern_name(pattern_id), iteration, seed)
if os.path.exists(pattern_iter_output_dir) and not overwrite_dir:
logger.warning(
f"Path {pattern_iter_output_dir} already exists, skipping it..."
)
continue
if not os.path.exists(pattern_iter_output_dir) and local_rank in [
-1, 0
]:
os.makedirs(pattern_iter_output_dir)
wrapper = init_model(model_config)
# Training
if do_train:
if ipet_data_dir:
p = os.path.join(
ipet_data_dir,
"{}-i{}-train.bin".format(pattern_name(pattern_id),
iteration))
ipet_train_data = InputExample.load_examples(p)
for example in ipet_train_data:
example.logits = None
else:
ipet_train_data = None
results_dict.update(
train_single_model(
wrapper,
train_data,
train_config,
pattern_iter_output_dir,
dev_data,
eval_config,
ipet_train_data=ipet_train_data,
unlabeled_data=unlabeled_data,
return_train_set_results=False,
local_rank=local_rank,
))
with open(os.path.join(pattern_iter_output_dir, "results.txt"),
"w") as fh:
fh.write(str(results_dict))
if local_rank in [-1, 0]:
logger.info("Saving trained model at {}...".format(
pattern_iter_output_dir))
train_config.save(
os.path.join(pattern_iter_output_dir,
"train_config.json"))
eval_config.save(
os.path.join(pattern_iter_output_dir,
"eval_config.json"))
logger.info("Saving complete")
if save_unlabeled_logits:
logits = evaluate(wrapper,
unlabeled_data,
eval_config,
local_rank=local_rank)["logits"]
save_logits(
os.path.join(pattern_iter_output_dir,
"logits.txt"), logits)
if not do_eval:
wrapper.model = None
wrapper = None
torch.cuda.empty_cache()
# Evaluation
if do_eval:
logger.info("Starting evaluation...")
try:
wrapper = TransformerModelWrapper.from_pretrained(
pattern_iter_output_dir)
except OSError:
warnings.warn(
"No model found saved, proceeding with current model instead of best"
)
pass
for split, eval_data in {
"dev": dev_data,
"test": test_data
}.items():
if eval_data is None:
continue
eval_result = evaluate(wrapper,
eval_data,
eval_config,
priming_data=train_data,
local_rank=local_rank)
if local_rank in [-1, 0]:
save_predictions(
os.path.join(pattern_iter_output_dir,
"predictions.jsonl"), wrapper,
eval_result)
save_logits(
os.path.join(pattern_iter_output_dir,
"eval_logits.txt"),
eval_result["logits"])
scores = eval_result["scores"]
logger.info(
"--- {} result (pattern_id={}, iteration={}) ---".
format(split, pattern_id, iteration))
logger.info(scores)
results_dict[f"{split}_set_after_training"] = scores
with open(
os.path.join(pattern_iter_output_dir,
"results.json"), "w") as fh:
json.dump(results_dict, fh)
for metric, value in scores.items():
results[split][metric][pattern_id].append(value)
wrapper.model = None
wrapper = None
torch.cuda.empty_cache()
if do_eval:
logger.info("=== OVERALL RESULTS ===")
results_to_log = _write_results(
os.path.join(output_dir, "result_test.txt"), results)
else:
logger.info("=== ENSEMBLE TRAINING COMPLETE ===")
results_to_log = None
if do_train and not save_model:
outputs = os.listdir(pattern_iter_output_dir)
for item in outputs:
if item.endswith(".bin"):
os.remove(os.path.join(pattern_iter_output_dir, item))
return results_to_log
def main():
parser = argparse.ArgumentParser()
# required parameters
parser.add_argument(
"--output_dir",
default=None,
type=str,
required=True,
help=
"The output directory. The verbalizers are written to a file 'verbalizer.json' in this directory.",
)
parser.add_argument(
"--data_dir",
default=None,
type=str,
required=True,
help="The input data dir. Should contain the data files for the task.",
)
parser.add_argument(
"--model_type",
default=None,
type=str,
required=True,
help="The model type",
)
parser.add_argument(
"--model_name_or_path",
default=None,
type=str,
required=True,
help="Path to pre-trained model or shortcut name",
)
parser.add_argument(
"--task_name",
default=None,
type=str,
required=True,
help="The name of the task to train selected in the list: " +
", ".join(PROCESSORS.keys()),
)
# verbalizer search hyperparameters
parser.add_argument(
"--normalize",
action="store_true",
help=
"Whether to normalize the loss as proposed in the paper. It is recommended to set this to 'true'.",
)
parser.add_argument(
"--combine_patterns",
action="store_true",
help=
"If set to true, a single joint verbalizer is searched for all patterns",
)
parser.add_argument(
"--num_candidates",
default=1000,
type=int,
help=
"The number of candidate tokens to consider as verbalizers (see Section 4.1 of the paper)",
)
parser.add_argument(
"--words_per_label",
default=10,
type=int,
help="The number of verbalizer tokens to assign to each label",
)
parser.add_argument(
"--score_fct",
default="llr",
choices=["llr", "ce", "random"],
help=
"The function used to score verbalizers. Choices are: the log-likelihood ratio loss proposed in the paper "
"('llr'), cross-entropy loss ('ce') and 'random', which assigns random tokens to each label.",
)
# other optional parameters
parser.add_argument(
"--train_examples",
default=50,
type=int,
help=
"The total number of train examples to use, where -1 equals all examples.",
)
parser.add_argument(
"--pattern_ids",
default=[0],
type=int,
nargs="+",
help="The ids of the PVPs to be used",
)
parser.add_argument(
"--max_seq_length",
default=256,
type=int,
help=
"The maximum total input sequence length after tokenization. Sequences longer "
"than this will be truncated, sequences shorter will be padded.",
)
parser.add_argument(
"--per_gpu_eval_batch_size",
default=8,
type=int,
help="Batch size per GPU/CPU for evaluation.",
)
parser.add_argument(
"--words_file",
default=None,
type=str,
help=
"Path to a file containing (unlabeled) texts from the task's domain. This text is used to compute "
"verbalization candidates by selecting the most frequent words.",
)
parser.add_argument(
"--max_words",
default=10000,
type=int,
help=
"Only the 10,000 tokens that occur most frequently in the task’s unlabeled data (see --words_file) are "
"considered as verbalization candidates",
)
parser.add_argument(
"--additional_input_examples",
type=str,
help=
"An optional path to an additional set of input examples (e.g., obtained using iPET)",
)
parser.add_argument("--seed",
default=42,
type=int,
help="random seed for initialization")
args = parser.parse_args()
random.seed(args.seed)
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
with open(os.path.join(args.output_dir, "config.txt"),
"w",
encoding="utf8") as fh:
json.dump(args.__dict__, fh, indent=2)
# setup gpu/cpu
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
args.n_gpu = torch.cuda.device_count()
# prepare task
args.task_name = args.task_name.lower()
if args.task_name not in PROCESSORS:
raise ValueError("Task not found: {}".format(args.task_name))
processor = PROCESSORS[args.task_name]()
args.label_list = processor.get_labels()
args.cache_dir = ""
args.do_lower_case = False
args.verbalizer_file = None
args.wrapper_type = "mlm"
# get training data
train_examples_per_label = (eq_div(args.train_examples, len(
args.label_list)) if args.train_examples != -1 else -1)
train_data = load_examples(
args.task_name,
args.data_dir,
set_type=TRAIN_SET,
num_examples_per_label=train_examples_per_label,
)
if args.additional_input_examples:
additional_data = InputExample.load_examples(
args.additional_input_examples)
train_data += additional_data
logger.info(
f"Loaded {len(additional_data)} additional examples from {args.additional_input_examples}, total"
f"training set size is now {len(train_data)}")
expected = {
label: np.array([1 if x.label == label else 0 for x in train_data])
for label in args.label_list
}
if args.words_file:
with open(args.words_file, "r", encoding="utf8") as fh:
word_counts = Counter(fh.read().split())
else:
word_counts = None
tokenizer_class = MODEL_CLASSES[args.model_type]["tokenizer"]
tokenizer = tokenizer_class.from_pretrained(args.model_name_or_path)
word2idx = get_word_to_id_map(tokenizer,
word_counts=word_counts,
max_words=args.max_words)
logits = []
for pattern_id in args.pattern_ids:
logger.info(f"Processing examples with pattern id {pattern_id}...")
args.pattern_id = pattern_id
config = WrapperConfig(
model_type=args.model_type,
model_name_or_path=args.model_name_or_path,
wrapper_type="mlm",
task_name=args.task_name,
max_seq_length=args.max_seq_length,
label_list=args.label_list,
pattern_id=args.pattern_id,
)
wrapper = TransformerModelWrapper(config)
wrapper.model.to(device)
# modify all patterns so that they return a single text segment instead of two segments
get_parts = wrapper.preprocessor.pvp.get_parts
wrapper.preprocessor.pvp.get_parts = lambda example: (
get_parts(example)[0] + get_parts(example)[1],
[],
)
wrapper.preprocessor.pvp.convert_mlm_logits_to_cls_logits = lambda mask, x, _=None: x[
mask >= 0]
pattern_logits = wrapper.eval(
train_data,
device,
per_gpu_eval_batch_size=args.per_gpu_eval_batch_size,
n_gpu=args.n_gpu,
)["logits"]
pattern_logits = pattern_logits - np.expand_dims(
np.max(pattern_logits, axis=1), axis=1)
logits.append(pattern_logits)
logger.info("Starting verbalizer search...")
if args.combine_patterns:
avs = AutomaticVerbalizerSearch(word2idx, args.label_list, logits,
expected)
verbalizer = avs.find_verbalizer(
num_candidates=args.num_candidates,
words_per_label=args.words_per_label,
normalize=args.normalize,
score_fct=args.score_fct,
)
verbalizers = {
pattern_id: verbalizer
for pattern_id in args.pattern_ids
}
else:
verbalizers = {}
for idx, pattern_id in enumerate(args.pattern_ids):
avs = AutomaticVerbalizerSearch(word2idx, args.label_list,
[logits[idx]], expected)
verbalizers[pattern_id] = avs.find_verbalizer(
num_candidates=args.num_candidates,
words_per_label=args.words_per_label,
normalize=args.normalize,
score_fct=args.score_fct,
)
print(json.dumps(verbalizers, indent=2))
logger.info("Verbalizer search complete, writing output...")
with open(os.path.join(args.output_dir, "verbalizers.json"),
"w",
encoding="utf8") as fh:
json.dump(verbalizers, fh, indent=2)
logger.info("Done")
