utils.py (24804B)
1 # Copyright 2020 The HuggingFace Team. All rights reserved. 2 # 3 # Licensed under the Apache License, Version 2.0 (the "License"); 4 # you may not use this file except in compliance with the License. 5 # You may obtain a copy of the License at 6 # 7 # http://www.apache.org/licenses/LICENSE-2.0 8 # 9 # Unless required by applicable law or agreed to in writing, software 10 # distributed under the License is distributed on an "AS IS" BASIS, 11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 # See the License for the specific language governing permissions and 13 # limitations under the License. 14 15 import itertools 16 import json 17 import linecache 18 import math 19 import os 20 import pickle 21 import socket 22 from logging import getLogger 23 from pathlib import Path 24 from typing import Callable, Dict, Iterable, List, Tuple, Union 25 26 import git 27 import numpy as np 28 import torch 29 import torch.distributed as dist 30 from rouge_score import rouge_scorer, scoring 31 from sacrebleu import corpus_bleu 32 from torch import nn 33 from torch.utils.data import Dataset, Sampler 34 35 from seq2seq.sentence_splitter import add_newline_to_end_of_each_sentence 36 from transformers import BartTokenizer, EvalPrediction, PreTrainedTokenizer 37 from transformers.file_utils import cached_property 38 39 try: 40 from fairseq.data.data_utils import batch_by_size 41 42 FAIRSEQ_AVAILABLE = True 43 except (ImportError, ModuleNotFoundError): 44 FAIRSEQ_AVAILABLE = False 45 46 47 def label_smoothed_nll_loss(lprobs, target, epsilon, ignore_index=-100): 48 """From fairseq""" 49 if target.dim() == lprobs.dim() - 1: 50 target = target.unsqueeze(-1) 51 nll_loss = -lprobs.gather(dim=-1, index=target) 52 smooth_loss = -lprobs.sum(dim=-1, keepdim=True) 53 if ignore_index is not None: 54 pad_mask = target.eq(ignore_index) 55 nll_loss.masked_fill_(pad_mask, 0.0) 56 smooth_loss.masked_fill_(pad_mask, 0.0) 57 else: 58 nll_loss = nll_loss.squeeze(-1) 59 smooth_loss = smooth_loss.squeeze(-1) 60 61 nll_loss = nll_loss.sum() # mean()? Scared to break other math. 62 smooth_loss = smooth_loss.sum() 63 eps_i = epsilon / lprobs.size(-1) 64 loss = (1.0 - epsilon) * nll_loss + eps_i * smooth_loss 65 return loss, nll_loss 66 67 68 def lmap(f: Callable, x: Iterable) -> List: 69 """list(map(f, x))""" 70 return list(map(f, x)) 71 72 73 def calculate_bleu(output_lns, refs_lns, **kwargs) -> dict: 74 """Uses sacrebleu's corpus_bleu implementation.""" 75 return {"bleu": round(corpus_bleu(output_lns, [refs_lns], **kwargs).score, 4)} 76 77 78 def build_compute_metrics_fn(task_name: str, tokenizer: PreTrainedTokenizer) -> Callable[[EvalPrediction], Dict]: 79 def non_pad_len(tokens: np.ndarray) -> int: 80 return np.count_nonzero(tokens != tokenizer.pad_token_id) 81 82 def decode_pred(pred: EvalPrediction) -> Tuple[List[str], List[str]]: 83 pred_str = tokenizer.batch_decode(pred.predictions, skip_special_tokens=True) 84 label_str = tokenizer.batch_decode(pred.label_ids, skip_special_tokens=True) 85 pred_str = lmap(str.strip, pred_str) 86 label_str = lmap(str.strip, label_str) 87 return pred_str, label_str 88 89 def summarization_metrics(pred: EvalPrediction) -> Dict: 90 pred_str, label_str = decode_pred(pred) 91 rouge: Dict = calculate_rouge(pred_str, label_str) 92 summ_len = np.round(np.mean(lmap(non_pad_len, pred.predictions)), 1) 93 rouge.update({"gen_len": summ_len}) 94 return rouge 95 96 def translation_metrics(pred: EvalPrediction) -> Dict: 97 pred_str, label_str = decode_pred(pred) 98 bleu: Dict = calculate_bleu(pred_str, label_str) 99 gen_len = np.round(np.mean(lmap(non_pad_len, pred.predictions)), 1) 100 bleu.update({"gen_len": gen_len}) 101 return bleu 102 103 compute_metrics_fn = summarization_metrics if "summarization" in task_name else translation_metrics 104 return compute_metrics_fn 105 106 107 def trim_batch( 108 input_ids, 109 pad_token_id, 110 attention_mask=None, 111 ): 112 """Remove columns that are populated exclusively by pad_token_id""" 113 keep_column_mask = input_ids.ne(pad_token_id).any(dim=0) 114 if attention_mask is None: 115 return input_ids[:, keep_column_mask] 116 else: 117 return (input_ids[:, keep_column_mask], attention_mask[:, keep_column_mask]) 118 119 120 class AbstractSeq2SeqDataset(Dataset): 121 def __init__( 122 self, 123 tokenizer, 124 data_dir, 125 max_source_length, 126 max_target_length, 127 type_path="train", 128 n_obs=None, 129 prefix="", 130 **dataset_kwargs 131 ): 132 super().__init__() 133 self.src_file = Path(data_dir).joinpath(type_path + ".source") 134 self.tgt_file = Path(data_dir).joinpath(type_path + ".target") 135 self.len_file = Path(data_dir).joinpath(type_path + ".len") 136 if os.path.exists(self.len_file): 137 self.src_lens = pickle_load(self.len_file) 138 self.used_char_len = False 139 else: 140 self.src_lens = self.get_char_lens(self.src_file) 141 self.used_char_len = True 142 self.max_source_length = max_source_length 143 self.max_target_length = max_target_length 144 assert min(self.src_lens) > 0, f"found empty line in {self.src_file}" 145 self.tokenizer = tokenizer 146 self.prefix = prefix if prefix is not None else "" 147 148 if n_obs is not None: 149 self.src_lens = self.src_lens[:n_obs] 150 self.pad_token_id = self.tokenizer.pad_token_id 151 self.dataset_kwargs = dataset_kwargs 152 dataset_kwargs.update({"add_prefix_space": True} if isinstance(self.tokenizer, BartTokenizer) else {}) 153 154 def __len__(self): 155 return len(self.src_lens) 156 157 @staticmethod 158 def get_char_lens(data_file): 159 return [len(x) for x in Path(data_file).open().readlines()] 160 161 @cached_property 162 def tgt_lens(self): 163 """Length in characters of target documents""" 164 return self.get_char_lens(self.tgt_file) 165 166 def make_sortish_sampler(self, batch_size, distributed=False, shuffle=True, **kwargs): 167 if distributed: 168 return DistributedSortishSampler(self, batch_size, shuffle=shuffle, **kwargs) 169 else: 170 return SortishSampler(self.src_lens, batch_size, shuffle=shuffle) 171 172 def make_dynamic_sampler(self, max_tokens_per_batch=1024, **kwargs): 173 assert FAIRSEQ_AVAILABLE, "Dynamic batch size requires `pip install fairseq`" 174 assert not self.used_char_len, "You must call python make_len_file.py before calling make_dynamic_sampler" 175 sorted_indices = list(self.make_sortish_sampler(1024, shuffle=False)) 176 177 def num_tokens_in_example(i): 178 return min(self.src_lens[i], self.max_target_length) 179 180 # call fairseq cython function 181 batch_sampler: List[List[int]] = batch_by_size( 182 sorted_indices, 183 num_tokens_fn=num_tokens_in_example, 184 max_tokens=max_tokens_per_batch, 185 required_batch_size_multiple=64, 186 ) 187 shuffled_batches = [batch_sampler[i] for i in np.random.permutation(range(len(batch_sampler)))] 188 # move the largest batch to the front to OOM quickly (uses an approximation for padding) 189 approximate_toks_per_batch = [max(self.src_lens[i] for i in batch) * len(batch) for batch in shuffled_batches] 190 largest_batch_idx = np.argmax(approximate_toks_per_batch) 191 shuffled_batches[0], shuffled_batches[largest_batch_idx] = ( 192 shuffled_batches[largest_batch_idx], 193 shuffled_batches[0], 194 ) 195 return shuffled_batches 196 197 def __getitem__(self, item): 198 raise NotImplementedError("You must implement this") 199 200 def collate_fn(self, batch): 201 raise NotImplementedError("You must implement this") 202 203 204 class LegacySeq2SeqDataset(AbstractSeq2SeqDataset): 205 def __getitem__(self, index) -> Dict[str, torch.Tensor]: 206 """Call tokenizer on src and tgt_lines""" 207 index = index + 1 # linecache starts at 1 208 source_line = self.prefix + linecache.getline(str(self.src_file), index).rstrip("\n") 209 tgt_line = linecache.getline(str(self.tgt_file), index).rstrip("\n") 210 assert source_line, f"empty source line for index {index}" 211 assert tgt_line, f"empty tgt line for index {index}" 212 source_inputs = self.encode_line(self.tokenizer, source_line, self.max_source_length) 213 target_inputs = self.encode_line(self.tokenizer, tgt_line, self.max_target_length) 214 215 source_ids = source_inputs["input_ids"].squeeze() 216 target_ids = target_inputs["input_ids"].squeeze() 217 src_mask = source_inputs["attention_mask"].squeeze() 218 return { 219 "input_ids": source_ids, 220 "attention_mask": src_mask, 221 "labels": target_ids, 222 } 223 224 def encode_line(self, tokenizer, line, max_length, pad_to_max_length=True, return_tensors="pt"): 225 """Only used by LegacyDataset""" 226 return tokenizer( 227 [line], 228 max_length=max_length, 229 padding="max_length" if pad_to_max_length else None, 230 truncation=True, 231 return_tensors=return_tensors, 232 **self.dataset_kwargs, 233 ) 234 235 def collate_fn(self, batch) -> Dict[str, torch.Tensor]: 236 input_ids = torch.stack([x["input_ids"] for x in batch]) 237 masks = torch.stack([x["attention_mask"] for x in batch]) 238 target_ids = torch.stack([x["labels"] for x in batch]) 239 pad_token_id = self.pad_token_id 240 y = trim_batch(target_ids, pad_token_id) 241 source_ids, source_mask = trim_batch(input_ids, pad_token_id, attention_mask=masks) 242 batch = { 243 "input_ids": source_ids, 244 "attention_mask": source_mask, 245 "labels": y, 246 } 247 return batch 248 249 250 class Seq2SeqDataset(AbstractSeq2SeqDataset): 251 """A dataset that calls prepare_seq2seq_batch.""" 252 253 def __getitem__(self, index) -> Dict[str, str]: 254 index = index + 1 # linecache starts at 1 255 source_line = self.prefix + linecache.getline(str(self.src_file), index).rstrip("\n") 256 tgt_line = linecache.getline(str(self.tgt_file), index).rstrip("\n") 257 assert source_line, f"empty source line for index {index}" 258 assert tgt_line, f"empty tgt line for index {index}" 259 return {"tgt_texts": tgt_line, "src_texts": source_line, "id": index - 1} 260 261 def collate_fn(self, batch) -> Dict[str, torch.Tensor]: 262 """Call prepare_seq2seq_batch.""" 263 batch_encoding: Dict[str, torch.Tensor] = self.tokenizer.prepare_seq2seq_batch( 264 [x["src_texts"] for x in batch], 265 tgt_texts=[x["tgt_texts"] for x in batch], 266 max_length=self.max_source_length, 267 max_target_length=self.max_target_length, 268 return_tensors="pt", 269 **self.dataset_kwargs, 270 ).data 271 batch_encoding["ids"] = torch.tensor([x["id"] for x in batch]) 272 return batch_encoding 273 274 275 class Seq2SeqDataCollator: 276 def __init__(self, tokenizer, data_args, tpu_num_cores=None): 277 self.tokenizer = tokenizer 278 self.pad_token_id = tokenizer.pad_token_id 279 assert ( 280 self.pad_token_id is not None 281 ), f"pad_token_id is not defined for ({self.tokenizer.__class__.__name__}), it must be defined." 282 self.data_args = data_args 283 self.tpu_num_cores = tpu_num_cores 284 self.dataset_kwargs = {"add_prefix_space": True} if isinstance(tokenizer, BartTokenizer) else {} 285 if data_args.src_lang is not None: 286 self.dataset_kwargs["src_lang"] = data_args.src_lang 287 if data_args.tgt_lang is not None: 288 self.dataset_kwargs["tgt_lang"] = data_args.tgt_lang 289 290 def __call__(self, batch) -> Dict[str, torch.Tensor]: 291 if hasattr(self.tokenizer, "prepare_seq2seq_batch"): 292 batch = self._encode(batch) 293 input_ids, attention_mask, labels = ( 294 batch["input_ids"], 295 batch["attention_mask"], 296 batch["labels"], 297 ) 298 else: 299 input_ids = torch.stack([x["input_ids"] for x in batch]) 300 attention_mask = torch.stack([x["attention_mask"] for x in batch]) 301 labels = torch.stack([x["labels"] for x in batch]) 302 303 labels = trim_batch(labels, self.pad_token_id) 304 input_ids, attention_mask = trim_batch(input_ids, self.pad_token_id, attention_mask=attention_mask) 305 306 batch = { 307 "input_ids": input_ids, 308 "attention_mask": attention_mask, 309 "labels": labels, 310 } 311 return batch 312 313 def _shift_right_t5(self, input_ids): 314 # shift inputs to the right 315 shifted_input_ids = input_ids.new_zeros(input_ids.shape) 316 shifted_input_ids[..., 1:] = input_ids[..., :-1].clone() 317 shifted_input_ids[..., 0] = self.pad_token_id 318 return shifted_input_ids 319 320 def _encode(self, batch) -> Dict[str, torch.Tensor]: 321 batch_encoding = self.tokenizer.prepare_seq2seq_batch( 322 [x["src_texts"] for x in batch], 323 tgt_texts=[x["tgt_texts"] for x in batch], 324 max_length=self.data_args.max_source_length, 325 max_target_length=self.data_args.max_target_length, 326 padding="max_length" if self.tpu_num_cores is not None else "longest", # TPU hack 327 return_tensors="pt", 328 **self.dataset_kwargs, 329 ) 330 return batch_encoding.data 331 332 333 class SortishSampler(Sampler): 334 "Go through the text data by order of src length with a bit of randomness. From fastai repo." 335 336 def __init__(self, data, batch_size, shuffle=True): 337 self.data, self.bs, self.shuffle = data, batch_size, shuffle 338 339 def __len__(self) -> int: 340 return len(self.data) 341 342 def __iter__(self): 343 return iter(sortish_sampler_indices(self.data, self.bs, shuffle=self.shuffle)) 344 345 346 def sortish_sampler_indices(data: List, bs: int, shuffle=True) -> np.array: 347 "Go through the text data by order of src length with a bit of randomness. From fastai repo." 348 if not shuffle: 349 return np.argsort(np.array(data) * -1) 350 351 def key_fn(i): 352 return data[i] 353 354 idxs = np.random.permutation(len(data)) 355 sz = bs * 50 356 ck_idx = [idxs[i: i + sz] for i in range(0, len(idxs), sz)] 357 sort_idx = np.concatenate([sorted(s, key=key_fn, reverse=True) for s in ck_idx]) 358 sz = bs 359 ck_idx = [sort_idx[i: i + sz] for i in range(0, len(sort_idx), sz)] 360 max_ck = np.argmax([key_fn(ck[0]) for ck in ck_idx]) # find the chunk with the largest key, 361 ck_idx[0], ck_idx[max_ck] = ck_idx[max_ck], ck_idx[0] # then make sure it goes first. 362 sort_idx = np.concatenate(np.random.permutation(ck_idx[1:])) if len(ck_idx) > 1 else np.array([], dtype=np.int) 363 sort_idx = np.concatenate((ck_idx[0], sort_idx)) 364 return sort_idx 365 366 367 class DistributedSortishSampler(Sampler): 368 """Copied from torch DistributedSampler""" 369 370 def __init__(self, dataset, batch_size, num_replicas=None, rank=None, add_extra_examples=True, shuffle=True): 371 if num_replicas is None: 372 if not dist.is_available(): 373 raise RuntimeError("Requires distributed package to be available") 374 num_replicas = dist.get_world_size() 375 if rank is None: 376 if not dist.is_available(): 377 raise RuntimeError("Requires distributed package to be available") 378 rank = dist.get_rank() 379 self.dataset = dataset 380 self.num_replicas = num_replicas 381 self.rank = rank 382 self.epoch = 0 383 if add_extra_examples: 384 self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.num_replicas)) 385 self.total_size = self.num_samples * self.num_replicas 386 else: 387 self.total_size = len(dataset) 388 self.num_samples = len(self.available_indices) 389 self.batch_size = batch_size 390 self.add_extra_examples = add_extra_examples 391 self.shuffle = shuffle 392 393 def __iter__(self) -> Iterable: 394 g = torch.Generator() 395 g.manual_seed(self.epoch) 396 397 sortish_data = [self.dataset.src_lens[i] for i in self.available_indices] 398 sortish_indices = sortish_sampler_indices(sortish_data, self.batch_size, shuffle=self.shuffle) 399 indices = [self.available_indices[i] for i in sortish_indices] 400 assert len(indices) == self.num_samples 401 return iter(indices) 402 403 @cached_property 404 def available_indices(self) -> np.array: 405 indices = list(range(len(self.dataset))) 406 # add extra samples to make it evenly divisible 407 indices += indices[: (self.total_size - len(indices))] 408 assert len(indices) == self.total_size 409 # subsample 410 available_indices = indices[self.rank: self.total_size: self.num_replicas] 411 return available_indices 412 413 def __len__(self): 414 return self.num_samples 415 416 def set_epoch(self, epoch): 417 self.epoch = epoch 418 419 420 logger = getLogger(__name__) 421 422 423 def use_task_specific_params(model, task): 424 """Update config with summarization specific params.""" 425 task_specific_params = model.config.task_specific_params 426 427 if task_specific_params is not None: 428 pars = task_specific_params.get(task, {}) 429 logger.info(f"setting model.config to task specific params for {task}:\n {pars}") 430 logger.info("note: command line args may override some of these") 431 model.config.update(pars) 432 433 434 def pickle_load(path): 435 """pickle.load(path)""" 436 with open(path, "rb") as f: 437 return pickle.load(f) 438 439 440 def pickle_save(obj, path): 441 """pickle.dump(obj, path)""" 442 with open(path, "wb") as f: 443 return pickle.dump(obj, f) 444 445 446 def flatten_list(summary_ids: List[List]): 447 return [x for x in itertools.chain.from_iterable(summary_ids)] 448 449 450 def save_git_info(folder_path: str) -> None: 451 """Save git information to output_dir/git_log.json""" 452 repo_infos = get_git_info() 453 save_json(repo_infos, os.path.join(folder_path, "git_log.json")) 454 455 456 def save_json(content, path, indent=4, **json_dump_kwargs): 457 with open(path, "w") as f: 458 json.dump(content, f, indent=indent, sort_keys=True, **json_dump_kwargs) 459 460 461 def load_json(path): 462 with open(path) as f: 463 return json.load(f) 464 465 466 def get_git_info(): 467 try: 468 repo = git.Repo(search_parent_directories=True) 469 repo_infos = { 470 "repo_id": str(repo), 471 "repo_sha": str(repo.head.object.hexsha), 472 "repo_branch": str(repo.active_branch), 473 "hostname": str(socket.gethostname()), 474 } 475 return repo_infos 476 except TypeError: 477 return { 478 "repo_id": None, 479 "repo_sha": None, 480 "repo_branch": None, 481 "hostname": None, 482 } 483 484 485 ROUGE_KEYS = ["rouge1", "rouge2", "rougeL", "rougeLsum"] 486 487 488 def extract_rouge_mid_statistics(dct): 489 new_dict = {} 490 for k1, v1 in dct.items(): 491 mid = v1.mid 492 new_dict[k1] = {stat: round(getattr(mid, stat), 4) for stat in ["precision", "recall", "fmeasure"]} 493 return new_dict 494 495 496 def calculate_rouge( 497 pred_lns: List[str], 498 tgt_lns: List[str], 499 use_stemmer=True, 500 rouge_keys=ROUGE_KEYS, 501 return_precision_and_recall=False, 502 bootstrap_aggregation=True, 503 newline_sep=True, 504 ) -> Dict: 505 """Calculate rouge using rouge_scorer package. 506 Args: 507 pred_lns: list of summaries generated by model 508 tgt_lns: list of groundtruth summaries (e.g. contents of val.target) 509 use_stemmer: Bool indicating whether Porter stemmer should be used to 510 strip word suffixes to improve matching. 511 rouge_keys: which metrics to compute, defaults to rouge1, rouge2, rougeL, rougeLsum 512 return_precision_and_recall: (False) whether to also return precision and recall. 513 bootstrap_aggregation: whether to do the typical bootstrap resampling of scores. Defaults to True, if False 514 this function returns a collections.defaultdict[metric: list of values for each observation for each subscore]`` 515 newline_sep:(default=True) whether to add newline between sentences. This is essential for calculation rougeL 516 on multi sentence summaries (CNN/DM dataset). 517 Returns: 518 Dict[score: value] if aggregate else defaultdict(list) keyed by rouge_keys 519 """ 520 scorer = rouge_scorer.RougeScorer(rouge_keys, use_stemmer=use_stemmer) 521 aggregator = scoring.BootstrapAggregator() 522 for pred, tgt in zip(tgt_lns, pred_lns): 523 # rougeLsum expects "\n" separated sentences within a summary 524 if newline_sep: 525 pred = add_newline_to_end_of_each_sentence(pred) 526 tgt = add_newline_to_end_of_each_sentence(tgt) 527 scores = scorer.score(pred, tgt) 528 aggregator.add_scores(scores) 529 530 if bootstrap_aggregation: 531 result = aggregator.aggregate() 532 if return_precision_and_recall: 533 return extract_rouge_mid_statistics(result) # here we return dict 534 else: 535 return {k: round(v.mid.fmeasure * 100, 4) for k, v in result.items()} 536 537 else: 538 return aggregator._scores # here we return defaultdict(list) 539 540 541 # Utilities for freezing parameters and checking whether they are frozen 542 543 544 def freeze_params(model: nn.Module): 545 """Set requires_grad=False for each of model.parameters()""" 546 for par in model.parameters(): 547 par.requires_grad = False 548 549 550 def freeze_embeds(model): 551 """Freeze token embeddings and positional embeddings for bart, just token embeddings for t5.""" 552 model_type = model.config.model_type 553 554 if model_type == "t5": 555 freeze_params(model.shared) 556 for d in [model.encoder, model.decoder]: 557 freeze_params(d.embed_tokens) 558 elif model_type == "fsmt": 559 for d in [model.model.encoder, model.model.decoder]: 560 freeze_params(d.embed_positions) 561 freeze_params(d.embed_tokens) 562 else: 563 freeze_params(model.model.shared) 564 for d in [model.model.encoder, model.model.decoder]: 565 freeze_params(d.embed_positions) 566 freeze_params(d.embed_tokens) 567 568 569 def grad_status(model: nn.Module) -> Iterable: 570 return (par.requires_grad for par in model.parameters()) 571 572 573 def any_requires_grad(model: nn.Module) -> bool: 574 return any(grad_status(model)) 575 576 577 def assert_all_frozen(model): 578 model_grads: List[bool] = list(grad_status(model)) 579 n_require_grad = sum(lmap(int, model_grads)) 580 npars = len(model_grads) 581 assert not any(model_grads), f"{n_require_grad / npars:.1%} of {npars} weights require grad" 582 583 584 def assert_not_all_frozen(model): 585 model_grads: List[bool] = list(grad_status(model)) 586 npars = len(model_grads) 587 assert any(model_grads), f"none of {npars} weights require grad" 588 589 590 def parse_numeric_n_bool_cl_kwargs(unparsed_args: List[str]) -> Dict[str, Union[int, float, bool]]: 591 """ 592 Parse an argv list of unspecified command line args to a dict. 593 Assumes all values are either numeric or boolean in the form of true/false. 594 """ 595 result = {} 596 assert len(unparsed_args) % 2 == 0, f"got odd number of unparsed args: {unparsed_args}" 597 num_pairs = len(unparsed_args) // 2 598 for pair_num in range(num_pairs): 599 i = 2 * pair_num 600 assert unparsed_args[i].startswith("--") 601 if unparsed_args[i + 1].lower() == "true": 602 value = True 603 elif unparsed_args[i + 1].lower() == "false": 604 value = False 605 else: 606 try: 607 value = int(unparsed_args[i + 1]) 608 except ValueError: 609 value = float(unparsed_args[i + 1]) # this can raise another informative ValueError 610 611 result[unparsed_args[i][2:]] = value 612 return result 613 614 615 def write_txt_file(ordered_tgt, path): 616 f = Path(path).open("w") 617 for ln in ordered_tgt: 618 f.write(ln + "\n") 619 f.flush() 620 621 622 def chunks(lst, n): 623 """Yield successive n-sized chunks from lst.""" 624 for i in range(0, len(lst), n): 625 yield lst[i: i + n] 626 627 628 def check_output_dir(args, expected_items=0): 629 """ 630 Checks whether to bail out if output_dir already exists and has more than expected_items in it 631 `args`: needs to have the following attributes of `args`: 632 - output_dir 633 - do_train 634 - overwrite_output_dir 635 `expected_items`: normally 0 (default) - i.e. empty dir, but in some cases a few files are expected (e.g. recovery from OOM) 636 """ 637 if ( 638 os.path.exists(args.output_dir) 639 and len(os.listdir(args.output_dir)) > expected_items 640 and args.do_train 641 and not args.overwrite_output_dir 642 ): 643 raise ValueError( 644 f"Output directory ({args.output_dir}) already exists and " 645 f"has {len(os.listdir(args.output_dir))} items in it (expected {expected_items} items). " 646 "Use --overwrite_output_dir to overcome." 647 )