File size: 13,027 Bytes
8c3d6d6 |
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 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 |
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_files", required=True, nargs='+', default=[])
parser.add_argument("--epochs", default=1, type=int)
parser.add_argument("--batch_size", default=32, type=int)
parser.add_argument("--max_source_length", default=320, type=int)
parser.add_argument("--max_target_length", default=64, type=int)
parser.add_argument("--name", required=True)
parser.add_argument("--train_size", default=10*1000*1000, type=int)
parser.add_argument("--eval_size", default=10000, type=int)
parser.add_argument("--fp16", 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):
print("open", self.filepath)
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):
return [raw_example[0], raw_example[1]]
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, filepaths, num_samples):
self.num_samples = num_samples
self.datasets = []
self.data_iterators = []
for filepath in filepaths:
if 'reddit_title_text' in filepath:
dataset = RedditTitleDataset(filepath)
if 'stackexchange_archive/jsonl' in filepath:
dataset = StackExchangeTitleBodyDataset(filepath)
else:
dataset = PairDataset(filepath)
self.datasets.append(dataset)
self.data_iterators.append(iter(dataset))
def __len__(self):
return self.num_samples
def __iter__(self):
while True:
for dataset in self.data_iterators:
yield next(dataset)
random.shuffle(self.data_iterators)
def delete_examples_cache(self):
for dataset in self.datasets:
dataset.examples = []
def main():
############ Model
model = AutoModelForSeq2SeqLM.from_pretrained(args.model_name)
tokenizer = AutoTokenizer.from_pretrained(args.model_name)
save_steps = 1000
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)
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_files, 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
print("Target:", eval_dataset[0][0])
print("Input:", eval_dataset[0][1])
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.py --model_name google/t5-v1_1-small --train_files /home/stackexchange_archive/jsonl/academia.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/android.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/anime.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/apple.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/arduino.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/askubuntu.com.jsonl.gz /home/stackexchange_archive/jsonl/astronomy.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/aviation.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/bicycles.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/biology.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/bitcoin.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/blender.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/boardgames.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/chemistry.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/christianity.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/civicrm.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/codereview.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/cooking.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/craftcms.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/crypto.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/cs.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/cstheory.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/datascience.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/dba.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/diy.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/drupal.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/dsp.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/economics.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/electronics.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/ell.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/emacs.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/engineering.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/english.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/ethereum.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/expressionengine.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/french.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/gamedev.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/gaming.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/gardening.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/german.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/gis.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/graphicdesign.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/hinduism.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/history.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/islam.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/japanese.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/judaism.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/law.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/magento.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/math.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/mathematica.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/mathoverflow.net.jsonl.gz /home/stackexchange_archive/jsonl/mechanics.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/meta.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/meta.stackoverflow.com.jsonl.gz /home/stackexchange_archive/jsonl/money.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/movies.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/music.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/networkengineering.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/philosophy.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/photo.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/physics.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/politics.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/puzzling.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/quant.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/raspberrypi.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/rpg.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/rus.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/salesforce.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/scifi.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/security.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/serverfault.com.jsonl.gz /home/stackexchange_archive/jsonl/sharepoint.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/skeptics.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/small_stackexchanges.jsonl.gz /home/stackexchange_archive/jsonl/softwareengineering.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/softwarerecs.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/space.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/stackoverflow.com-Posts.jsonl.gz /home/stackexchange_archive/jsonl/stats.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/superuser.com.jsonl.gz /home/stackexchange_archive/jsonl/tex.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/travel.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/unix.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/ux.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/vi.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/webapps.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/webmasters.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/wordpress.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/workplace.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/worldbuilding.stackexchange.com.jsonl.gz /home/stackexchange_archive/jsonl/writers.stackexchange.com.jsonl.gz --name stackexchange_title_text_all --train_size 100000000 --max_source_length 384 |