LegumMagister/MNLP_M3_document_encoder
083
1"""2Train script for a single file3 4Need to set the TPU address first:5export XRT_TPU_CONFIG="localservice;0;localhost:51011"6"""7 8import torch.multiprocessing as mp9import threading10import time11import random12import sys13import argparse14import gzip15import json16import logging17import tqdm18import torch19from torch import nn20from torch.utils.data import DataLoader21import torch22import torch_xla23import torch_xla.core24import torch_xla.core.functions25import torch_xla.core.xla_model as xm26import torch_xla.distributed.xla_multiprocessing as xmp27import torch_xla.distributed.parallel_loader as pl28import os29from shutil import copyfile30 31 32from transformers import (33 AdamW,34 AutoModel,35 AutoTokenizer,36 get_linear_schedule_with_warmup,37 set_seed,38)39 40class AutoModelForSentenceEmbedding(nn.Module):41 def __init__(self, model_name, tokenizer, normalize=True):42 super(AutoModelForSentenceEmbedding, self).__init__()43 44 self.model = AutoModel.from_pretrained(model_name)45 self.normalize = normalize46 self.tokenizer = tokenizer47 48 def forward(self, **kwargs):49 model_output = self.model(**kwargs)50 embeddings = self.mean_pooling(model_output, kwargs['attention_mask'])51 if self.normalize:52 embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)53 54 return embeddings55 56 def mean_pooling(self, model_output, attention_mask):57 token_embeddings = model_output[0] # First element of model_output contains all token embeddings58 input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()59 return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)60 61 def save_pretrained(self, output_path):62 if xm.is_master_ordinal():63 self.tokenizer.save_pretrained(output_path)64 self.model.config.save_pretrained(output_path)65 66 xm.save(self.model.state_dict(), os.path.join(output_path, "pytorch_model.bin"))67 68 69 70 71def train_function(index, args, queue):72 tokenizer = AutoTokenizer.from_pretrained(args.model)73 model = AutoModelForSentenceEmbedding(args.model, tokenizer)74 75 76 ### Train Loop77 device = xm.xla_device()78 model = model.to(device)79 80 # Instantiate optimizer81 optimizer = AdamW(params=model.parameters(), lr=2e-5, correct_bias=True)82 83 lr_scheduler = get_linear_schedule_with_warmup(84 optimizer=optimizer,85 num_warmup_steps=500,86 num_training_steps=args.steps,87 )88 89 # Now we train the model90 cross_entropy_loss = nn.CrossEntropyLoss()91 max_grad_norm = 192 93 model.train()94 95 for global_step in tqdm.trange(args.steps, disable=not xm.is_master_ordinal()):96 #### Get the batch data97 batch = queue.get()98 #print(index, "batch {}x{}".format(len(batch), ",".join([str(len(b)) for b in batch])))99 100 101 if len(batch[0]) == 2: #(anchor, positive)102 text1 = tokenizer([b[0] for b in batch], return_tensors="pt", max_length=args.max_length, truncation=True, padding="max_length")103 text2 = tokenizer([b[1] for b in batch], return_tensors="pt", max_length=args.max_length, truncation=True, padding="max_length")104 105 ### Compute embeddings106 embeddings_a = model(**text1.to(device))107 embeddings_b = model(**text2.to(device))108 109 ### Gather all embedings 110 embeddings_a = torch_xla.core.functions.all_gather(embeddings_a)111 embeddings_b = torch_xla.core.functions.all_gather(embeddings_b)112 113 ### Compute similarity scores 512 x 512114 scores = torch.mm(embeddings_a, embeddings_b.transpose(0, 1)) * args.scale115 116 ### Compute cross-entropy loss117 labels = torch.tensor(range(len(scores)), dtype=torch.long, device=embeddings_a.device) # Example a[i] should match with b[i]118 119 ## Symmetric loss as in CLIP120 loss = (cross_entropy_loss(scores, labels) + cross_entropy_loss(scores.transpose(0, 1), labels)) / 2121 122 else: #(anchor, positive, negative)123 text1 = tokenizer([b[0] for b in batch], return_tensors="pt", max_length=args.max_length, truncation=True, padding="max_length")124 text2 = tokenizer([b[1] for b in batch], return_tensors="pt", max_length=args.max_length, truncation=True, padding="max_length")125 text3 = tokenizer([b[2] for b in batch], return_tensors="pt", max_length=args.max_length, truncation=True, padding="max_length")126 127 embeddings_a = model(**text1.to(device))128 embeddings_b1 = model(**text2.to(device))129 embeddings_b2 = model(**text3.to(device))130 131 embeddings_a = torch_xla.core.functions.all_gather(embeddings_a)132 embeddings_b1 = torch_xla.core.functions.all_gather(embeddings_b1)133 embeddings_b2 = torch_xla.core.functions.all_gather(embeddings_b2)134 135 embeddings_b = torch.cat([embeddings_b1, embeddings_b2])136 137 ### Compute similarity scores 512 x 1024138 scores = torch.mm(embeddings_a, embeddings_b.transpose(0, 1)) * args.scale139 140 ### Compute cross-entropy loss141 labels = torch.tensor(range(len(scores)), dtype=torch.long, device=embeddings_a.device) # Example a[i] should match with b[i]142 143 ## One-way loss144 loss = cross_entropy_loss(scores, labels)145 146 147 # Backward pass148 optimizer.zero_grad()149 loss.backward()150 torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)151 152 xm.optimizer_step(optimizer, barrier=True)153 lr_scheduler.step()154 155 156 #Save model157 if (global_step+1) % args.save_steps == 0:158 output_path = os.path.join(args.output, str(global_step+1))159 xm.master_print("save model: "+output_path)160 model.save_pretrained(output_path)161 162 163 output_path = os.path.join(args.output, "final")164 xm.master_print("save model final: "+ output_path)165 model.save_pretrained(output_path)166 167 168def produce_data(args, queue, filepaths, dataset_indices):169 global_batch_size = args.batch_size*args.nprocs #Global batch size170 size_per_dataset = int(global_batch_size / args.datasets_per_batch) #How many datasets per batch171 num_same_dataset = int(size_per_dataset / args.batch_size)172 print("producer", "global_batch_size", global_batch_size)173 print("producer", "size_per_dataset", size_per_dataset)174 print("producer", "num_same_dataset", num_same_dataset)175 176 datasets = []177 for filepath in filepaths:178 if "reddit_" in filepath: #Special dataset class for Reddit files179 data_obj = RedditDataset(filepath)180 else:181 data_obj = Dataset(filepath)182 datasets.append(iter(data_obj)) 183 184 # Store if dataset is in a 2 col or 3 col format185 num_cols = {idx: len(next(dataset)) for idx, dataset in enumerate(datasets)}186 187 while True:188 texts_in_batch = set()189 batch_format = None #2 vs 3 col format for this batch190 191 #Add data from several sub datasets192 for _ in range(args.datasets_per_batch):193 valid_dataset = False #Check that datasets have the same 2/3 col format194 while not valid_dataset:195 data_idx = random.choice(dataset_indices)196 if batch_format is None:197 batch_format = num_cols[data_idx]198 valid_dataset = True199 else: #Check that this dataset has the same format200 valid_dataset = (batch_format == num_cols[data_idx])201 202 #Get data from this dataset203 dataset = datasets[data_idx]204 for _ in range(num_same_dataset):205 for _ in range(args.nprocs):206 batch_device = [] #A batch for one device207 while len(batch_device) < args.batch_size:208 sample = next(dataset)209 in_batch = False210 for text in sample:211 if text in texts_in_batch:212 in_batch = True213 break214 215 if not in_batch:216 for text in sample:217 texts_in_batch.add(text)218 batch_device.append(sample)219 220 queue.put(batch_device)221 222 223class RedditDataset:224 """225 A class that handles the reddit data files226 """227 def __init__(self, filepath):228 self.filepath = filepath229 230 def __iter__(self):231 while True:232 with gzip.open(self.filepath, "rt") as fIn:233 for line in fIn:234 data = json.loads(line)235 236 if "response" in data and "context" in data:237 yield [data["response"], data["context"]]238 239class Dataset:240 """241 A class that handles one dataset242 """243 def __init__(self, filepath):244 self.filepath = filepath245 246 def __iter__(self):247 max_dataset_size = 10*1000*1000 #Cache small datasets in memory248 dataset = []249 data_format = None250 251 while dataset is None or len(dataset) == 0:252 with gzip.open(self.filepath, "rt") as fIn:253 for line in fIn:254 data = json.loads(line)255 if isinstance(data, dict):256 data = data['texts']257 258 if data_format is None:259 data_format = len(data)260 261 #Ensure that all entries are of the same 2/3 col format262 assert len(data) == data_format263 264 if dataset is not None:265 dataset.append(data)266 if len(dataset) >= max_dataset_size:267 dataset = None268 269 yield data270 271 # Data loaded. Now stream to the queue272 # Shuffle for each epoch273 while True:274 random.shuffle(dataset)275 for data in dataset:276 yield data277 278 279 280if __name__ == "__main__":281 parser = argparse.ArgumentParser()282 parser.add_argument('--model', default='nreimers/MiniLM-L6-H384-uncased')283 parser.add_argument('--steps', type=int, default=2000)284 parser.add_argument('--save_steps', type=int, default=10000)285 parser.add_argument('--batch_size', type=int, default=64)286 parser.add_argument('--max_length', type=int, default=128)287 parser.add_argument('--nprocs', type=int, default=8)288 parser.add_argument('--datasets_per_batch', type=int, default=2, help="Number of datasets per batch")289 parser.add_argument('--scale', type=float, default=20, help="Use 20 for cossim, and 1 when you work with unnormalized embeddings with dot product")290 parser.add_argument('--data_folder', default="/data", help="Folder with your dataset files")291 parser.add_argument('data_config', help="A data_config.json file")292 parser.add_argument('output')293 args = parser.parse_args()294 295 # Ensure global batch size is divisble by data_sample_size296 assert (args.batch_size*args.nprocs) % args.datasets_per_batch == 0297 298 logging.info("Output: "+args.output)299 if os.path.exists(args.output):300 print("Output folder already exists.")301 input("Continue?")302 303 # Write train script to output path304 os.makedirs(args.output, exist_ok=True)305 306 data_config_path = os.path.join(args.output, 'data_config.json')307 copyfile(args.data_config, data_config_path)308 309 train_script_path = os.path.join(args.output, 'train_script.py')310 copyfile(__file__, train_script_path)311 with open(train_script_path, 'a') as fOut:312 fOut.write("\n\n# Script was called via:\n#python " + " ".join(sys.argv))313 314 315 316 #Load data config317 with open(args.data_config) as fIn:318 data_config = json.load(fIn)319 320 queue = mp.Queue(maxsize=100*args.nprocs)321 322 filepaths = []323 dataset_indices = []324 for idx, data in enumerate(data_config):325 filepaths.append(os.path.join(os.path.expanduser(args.data_folder), data['name']))326 dataset_indices.extend([idx]*data['weight'])327 328 # Start producer329 p = mp.Process(target=produce_data, args=(args, queue, filepaths, dataset_indices))330 p.start()331 332 # Run training333 print("Start processes:", args.nprocs)334 xmp.spawn(train_function, args=(args, queue), nprocs=args.nprocs, start_method='fork')335 print("Training done")336 print("It might be that not all processes exit automatically. In that case you must manually kill this process.")337 print("With 'pkill python' you can kill all remaining python processes")338 p.kill()339 exit()340 341 342 343# Script was called via:344#python train_many_data_files_v2.py --steps 1000000 --batch_size 128 --model nreimers/MiniLM-L6-H384-uncased train_data_configs/all_datasets_v4.json output/all_datasets_v4_MiniLM-L6-H384-uncased-batch128