import argparse import logging from torch.utils.data import Dataset, IterableDataset import gzip import json from transformers import Seq2SeqTrainer, AutoModelForSeq2SeqLM, AutoTokenizer, Seq2SeqTrainingArguments import sys from datetime import datetime import torch import random from shutil import copyfile import os import wandb import re logging.basicConfig( format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", handlers=[logging.StreamHandler(sys.stdout)], ) parser = argparse.ArgumentParser() parser.add_argument("--model_name", default="google/t5-v1_1-base") parser.add_argument("--train_file", required=True) parser.add_argument("--epochs", default=1, type=int) parser.add_argument("--batch_size", default=16, type=int) parser.add_argument("--max_source_length", default=384, type=int) parser.add_argument("--max_target_length", default=64, type=int) parser.add_argument("--name", required=True) parser.add_argument("--train_size", default=100*1000*1000, type=int) parser.add_argument("--eval_size", default=10000, type=int) parser.add_argument("--fp16", default=False, action='store_true') parser.add_argument("--no_prefix", default=False, action='store_true') args = parser.parse_args() wandb.init(project="doc2query", name=f"{args.name}-{args.model_name}") class PairDataset: def __init__(self, filepath): self.filepath = filepath self.examples = [] def __iter__(self): with gzip.open(self.filepath, 'rt') as fIn: for line in fIn: example = self.get_example(json.loads(line)) if example is not None: self.examples.append(example) yield example while True: random.shuffle(self.examples) for ex in self.examples: yield ex def get_example(self, raw_example): if isinstance(raw_example, dict): if 'set' in raw_example: example = random.sample(raw_example['set'], 2) elif 'query' in raw_example: example = [raw_example['query'], random.choice(raw_example['pos'])] else: raise ValueError("Unknown format: "+str(raw_example)) else: example = [raw_example[0], raw_example[1]] return example class RedditTitleDataset(PairDataset): def get_example(self, raw_example): return [self.clean_title(raw_example['title']), raw_example['body']] def clean_title(self, text): text = text.replace("&", "&").strip() if text.startswith("["): text = re.sub("^\[[a-zA-Z0-9]+\]", "", text).strip() if text.endswith("]"): text = re.sub("\[[a-zA-Z0-9\.]+\]$", "", text).strip() if text.startswith("/r"): text = re.sub("^/[a-zA-Z0-9/]+[;,: \-]+", "", text).strip() return text class StackExchangeTitleBodyDataset(PairDataset): def get_example(self, raw_example): return raw_example['texts'] class MultiDataset(IterableDataset): def __init__(self, train_config_path, num_samples): self.num_samples = num_samples with open(train_config_path) as fIn: train_config = json.load(fIn) self.categories = [] self.files = {} self.file2dataset = {} self.file2datasetIter = {} for prefix in train_config: self.categories.extend([prefix]*train_config[prefix]['weight']) self.files[prefix] = [] for filename, weight in train_config[prefix]['files'].items(): self.files[prefix].extend([filename]*weight) dataset = self.OpenDataset(filename) self.file2dataset[filename] = dataset self.file2datasetIter[filename] = iter(dataset) random.shuffle(self.files[prefix]) random.shuffle(self.categories) def OpenDataset(self, filepath): if 'reddit_title_text' in filepath: dataset = RedditTitleDataset(filepath) elif 'stackexchange_archive/jsonl' in filepath: dataset = StackExchangeTitleBodyDataset(filepath) else: dataset = PairDataset(filepath) return dataset def __len__(self): return self.num_samples def __iter__(self): while True: category = random.choice(self.categories) filepath = random.choice(self.files[category]) dataset = self.file2datasetIter[filepath] pair = next(dataset) #Add prefix to the input if not args.no_prefix: pair[1] = category+": "+pair[1].strip() yield pair def delete_examples_cache(self): for dataset in self.file2dataset.values(): dataset.examples = [] def main(): ############ Model model = AutoModelForSeq2SeqLM.from_pretrained(args.model_name) tokenizer = AutoTokenizer.from_pretrained(args.model_name) save_steps = 5000 output_dir = 'output/'+args.name+'-'+args.model_name.replace("/", "-")+'-'+datetime.now().strftime("%Y-%m-%d_%H-%M-%S") print("Output dir:", output_dir) # Write self to path os.makedirs(output_dir, exist_ok=True) copyfile(args.train_file, os.path.join(output_dir, 'data_config.json')) train_script_path = os.path.join(output_dir, 'train_script.py') copyfile(__file__, train_script_path) with open(train_script_path, 'a') as fOut: fOut.write("\n\n# Script was called via:\n#python " + " ".join(sys.argv)) #### training_args = Seq2SeqTrainingArguments( output_dir=output_dir, fp16=args.fp16, fp16_backend="amp", per_device_train_batch_size=args.batch_size, evaluation_strategy="steps", save_steps=save_steps, logging_steps=100, eval_steps=save_steps, #logging_steps, warmup_steps=1000, save_total_limit=1, num_train_epochs=args.epochs, report_to="wandb", ) ############ Arguments ############ Load datasets train_dataset = MultiDataset(args.train_file, args.train_size) train_dataset_iter = iter(train_dataset) eval_dataset = [next(train_dataset_iter) for _ in range(args.eval_size)] train_dataset.delete_examples_cache() #Make sure dev data is no re-used for training for i in range(50): print("Target:", eval_dataset[i][0]) print("Input:", eval_dataset[i][1]) print("\n\n===================\n\n") print("Train dataset len:", len(train_dataset)) def data_collator(examples): targets = [row[0] for row in examples] inputs = [row[1] for row in examples] label_pad_token_id = -100 model_inputs = tokenizer(inputs, max_length=args.max_source_length, padding=True, truncation=True, return_tensors='pt', pad_to_multiple_of=8 if training_args.fp16 else None) # Setup the tokenizer for targets with tokenizer.as_target_tokenizer(): labels = tokenizer(targets, max_length=args.max_target_length, padding=True, truncation=True, pad_to_multiple_of=8 if training_args.fp16 else None) # replace all tokenizer.pad_token_id in the labels by -100 to ignore padding in the loss. labels["input_ids"] = [ [(l if l != tokenizer.pad_token_id else label_pad_token_id) for l in label] for label in labels["input_ids"] ] model_inputs["labels"] = torch.tensor(labels["input_ids"]) return model_inputs ## Define the trainer trainer = Seq2SeqTrainer( model=model, args=training_args, train_dataset=train_dataset, eval_dataset=eval_dataset, tokenizer=tokenizer, data_collator=data_collator ) ### Save the model train_result = trainer.train() trainer.save_model() if __name__ == "__main__": main() # Script was called via: #python train_hf_trainer_prefix.py --train_file train_config.json --name all-datasets-v1-no_prefix --no_prefix