def main(): # See all possible arguments in src/transformers/training_args.py # or by passing the --help flag to this script. # We now keep distinct sets of args, for a cleaner separation of concerns. parser = HfArgumentParser( (ModelArguments, DataTrainingArguments, TrainingArguments)) if len(sys.argv) == 2 and sys.argv[1].endswith(".json"): # If we pass only one argument to the script and it's the path to a json file, # let's parse it to get our arguments. model_args, data_args, training_args = parser.parse_json_file( json_file=os.path.abspath(sys.argv[1])) else: model_args, data_args, training_args = parser.parse_args_into_dataclasses( ) if (os.path.exists(training_args.output_dir) and os.listdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir): raise ValueError( f"Output directory ({training_args.output_dir}) already exists and is not empty. Use --overwrite_output_dir to overcome." ) module = import_module("tasks") try: token_classification_task_clazz = getattr(module, model_args.task_type) token_classification_task: TokenClassificationTask = token_classification_task_clazz( ) except AttributeError: raise ValueError( f"Task {model_args.task_type} needs to be defined as a TokenClassificationTask subclass in {module}. " f"Available tasks classes are: {TokenClassificationTask.__subclasses__()}" ) # Setup logging logging.basicConfig( format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S", level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN, ) logger.warning( "Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s", training_args.local_rank, training_args.device, training_args.n_gpu, bool(training_args.local_rank != -1), training_args.fp16, ) # Set the verbosity to info of the Transformers logger (on main process only): if is_main_process(training_args.local_rank): transformers.utils.logging.set_verbosity_info() transformers.utils.logging.enable_default_handler() transformers.utils.logging.enable_explicit_format() logger.info("Training/evaluation parameters %s", training_args) # Set seed set_seed(training_args.seed) # Prepare CONLL-2003 task labels = token_classification_task.get_labels(data_args.labels) label_map: Dict[int, str] = {i: label for i, label in enumerate(labels)} num_labels = len(labels) # Load pretrained model and tokenizer # # Distributed training: # The .from_pretrained methods guarantee that only one local process can concurrently # download model & vocab. config = AutoConfig.from_pretrained( model_args.config_name if model_args.config_name else model_args.model_name_or_path, num_labels=num_labels, id2label=label_map, label2id={label: i for i, label in enumerate(labels)}, cache_dir=model_args.cache_dir, ) tokenizer = AutoTokenizer.from_pretrained( model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast, ) model = AutoModelForTokenClassification.from_pretrained( model_args.model_name_or_path, from_tf=bool(".ckpt" in model_args.model_name_or_path), config=config, cache_dir=model_args.cache_dir, ) # Get datasets train_dataset = (TokenClassificationDataset( token_classification_task=token_classification_task, data_dir=data_args.data_dir, tokenizer=tokenizer, labels=labels, model_type=config.model_type, max_seq_length=data_args.max_seq_length, overwrite_cache=data_args.overwrite_cache, mode=Split.train, ) if training_args.do_train else None) eval_dataset = (TokenClassificationDataset( token_classification_task=token_classification_task, data_dir=data_args.data_dir, tokenizer=tokenizer, labels=labels, model_type=config.model_type, max_seq_length=data_args.max_seq_length, overwrite_cache=data_args.overwrite_cache, mode=Split.dev, ) if training_args.do_eval else None) def align_predictions( predictions: np.ndarray, label_ids: np.ndarray) -> Tuple[List[int], List[int]]: preds = np.argmax(predictions, axis=2) batch_size, seq_len = preds.shape out_label_list = [[] for _ in range(batch_size)] preds_list = [[] for _ in range(batch_size)] for i in range(batch_size): for j in range(seq_len): if label_ids[i, j] != nn.CrossEntropyLoss().ignore_index: out_label_list[i].append(label_map[label_ids[i][j]]) preds_list[i].append(label_map[preds[i][j]]) return preds_list, out_label_list def compute_metrics(p: EvalPrediction) -> Dict: preds_list, out_label_list = align_predictions(p.predictions, p.label_ids) return { "accuracy_score": accuracy_score(out_label_list, preds_list), "precision": precision_score(out_label_list, preds_list), "recall": recall_score(out_label_list, preds_list), "f1": f1_score(out_label_list, preds_list), } # Data collator data_collator = DataCollatorWithPadding( tokenizer, pad_to_multiple_of=8) if training_args.fp16 else None # Initialize our Trainer trainer = Trainer( model=model, args=training_args, train_dataset=train_dataset, eval_dataset=eval_dataset, compute_metrics=compute_metrics, data_collator=data_collator, ) # Training if training_args.do_train: trainer.train(model_path=model_args.model_name_or_path if os.path. isdir(model_args.model_name_or_path) else None) trainer.save_model() # For convenience, we also re-save the tokenizer to the same directory, # so that you can share your model easily on huggingface.co/models =) if trainer.is_world_process_zero(): tokenizer.save_pretrained(training_args.output_dir) # Evaluation results = {} if training_args.do_eval: logger.info("*** Evaluate ***") result = trainer.evaluate() output_eval_file = os.path.join(training_args.output_dir, "eval_results.txt") if trainer.is_world_process_zero(): with open(output_eval_file, "w") as writer: logger.info("***** Eval results *****") for key, value in result.items(): logger.info(" %s = %s", key, value) writer.write("%s = %s\n" % (key, value)) results.update(result) # Predict if training_args.do_predict: test_dataset = TokenClassificationDataset( token_classification_task=token_classification_task, data_dir=data_args.data_dir, tokenizer=tokenizer, labels=labels, model_type=config.model_type, max_seq_length=data_args.max_seq_length, overwrite_cache=data_args.overwrite_cache, mode=Split.test, ) predictions, label_ids, metrics = trainer.predict(test_dataset) preds_list, _ = align_predictions(predictions, label_ids) output_test_results_file = os.path.join(training_args.output_dir, "test_results.txt") if trainer.is_world_process_zero(): with open(output_test_results_file, "w") as writer: for key, value in metrics.items(): logger.info(" %s = %s", key, value) writer.write("%s = %s\n" % (key, value)) # Save predictions output_test_predictions_file = os.path.join(training_args.output_dir, "test_predictions.txt") if trainer.is_world_process_zero(): with open(output_test_predictions_file, "w") as writer: with open(os.path.join(data_args.data_dir, "test.txt"), "r") as f: token_classification_task.write_predictions_to_file( writer, f, preds_list) return results
def main(): # See all possible arguments in src/transformers/training_args.py # or by passing the --help flag to this script. # We now keep distinct sets of args, for a cleaner separation of concerns. parser = HfArgumentParser( (ModelArguments, DataTrainingArguments, TrainingArguments)) if len(sys.argv) == 2 and sys.argv[1].endswith(".json"): # If we pass only one argument to the script and it's the path to a json file, # let's parse it to get our arguments. model_args, data_args, training_args = parser.parse_json_file( json_file=os.path.abspath(sys.argv[1])) else: model_args, data_args, training_args = parser.parse_args_into_dataclasses( ) if (os.path.exists(training_args.output_dir) and os.listdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir): raise ValueError( f"Output directory ({training_args.output_dir}) already exists and is not empty. Use --overwrite_output_dir to overcome." ) module = import_module("tasks") try: token_classification_task_clazz = getattr(module, model_args.task_type) token_classification_task: TokenClassificationTask = token_classification_task_clazz( ) except AttributeError: raise ValueError( f"Task {model_args.task_type} needs to be defined as a TokenClassificationTask subclass in {module}. " f"Available tasks classes are: {TokenClassificationTask.__subclasses__()}" ) # Setup logging if IS_IN_DOCKER_CONTAINER: log_level = logging.WARNING elif training_args.training_on_cloud: log_level = logging.WARNING else: log_level = logging.INFO if (training_args.local_rank in [-1, 0]) else logging.WARN logging.basicConfig( format= "%(asctime)s - %(levelname)s - %(name)s - %(lineno)d - %(message)s", datefmt="%m/%d/%Y %H:%M:%S", level=log_level) logger.warning(f"log level: {log_level}") logger.warning( "Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s", training_args.local_rank, training_args.device, training_args.n_gpu, bool(training_args.local_rank != -1), training_args.fp16, ) logger.info("Training/evaluation parameters:") for key, value in vars(training_args).items(): logger.info(" %s = %s", key, value) # Set seed set_seed(training_args.seed) # Prepare CONLL-2003 task labels = token_classification_task.get_labels(data_args.labels) label_map: Dict[int, str] = {i: label for i, label in enumerate(labels)} num_labels = len(labels) # Load pretrained model and tokenizer # # Distributed training: # The .from_pretrained methods guarantee that only one local process can concurrently # download model & vocab. config = AutoConfig.from_pretrained( model_args.config_name if model_args.config_name else model_args.model_name_or_path, num_labels=num_labels, id2label=label_map, label2id={label: i for i, label in enumerate(labels)}, cache_dir=model_args.cache_dir, gradient_checkpointing=training_args.gradient_checkpointing) tokenizer = AutoTokenizer.from_pretrained( model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast, additional_special_tokens=["[unused1]"]) model = AutoModelForTokenClassification.from_pretrained( model_args.model_name_or_path, from_tf=bool(".ckpt" in model_args.model_name_or_path), config=config, cache_dir=model_args.cache_dir, training_args=training_args) # Get datasets train_dataset = (TokenClassificationDataset( token_classification_task=token_classification_task, data_dir=data_args.data_dir, tokenizer=tokenizer, labels=labels, model_type=config.model_type, max_seq_length=data_args.max_seq_length, overwrite_cache=data_args.overwrite_cache, mode=Split.train, training_args=training_args, remove_label=Remove_Label) if training_args.do_train else None) eval_dataset = (TokenClassificationDataset( token_classification_task=token_classification_task, data_dir=data_args.data_dir, tokenizer=tokenizer, labels=labels, model_type=config.model_type, max_seq_length=data_args.max_seq_length, overwrite_cache=data_args.overwrite_cache, mode=Split.dev, training_args=training_args, remove_label=Remove_Label) if training_args.do_eval or training_args.do_predict_dev or training_args.evaluate_during_training else None) def align_predictions(predictions: np.ndarray, label_ids: np.ndarray, dateset=None) -> Tuple[List[int], List[int]]: if training_args.use_crf: mask = [] for feature in dateset.features: mask.append(feature.attention_mask) if mask: if isinstance(mask, list): mask = np.array(mask) mask = torch.from_numpy(mask).cuda() == 1 predictions = torch.from_numpy(predictions).cuda() preds = model.crf.decode(predictions, mask) else: preds = model.crf.decode(predictions) out_label_list = [] preds_list = [] for pred_one, label_one in zip(preds, label_ids.tolist()): out_label = [] pred = [] for p, l in zip(pred_one, label_one): out_label.append(label_map[l]) pred.append(label_map[p]) out_label = out_label[1:-1] pred = pred[1:-1] out_label_list.append(out_label) preds_list.append(pred) else: preds = np.argmax(predictions, axis=2) batch_size, seq_len = preds.shape out_label_list = [[] for _ in range(batch_size)] preds_list = [[] for _ in range(batch_size)] for i in range(batch_size): for j in range(seq_len): if label_ids[i, j] != nn.CrossEntropyLoss().ignore_index: out_label_list[i].append(label_map[label_ids[i][j]]) preds_list[i].append(label_map[preds[i][j]]) return preds_list, out_label_list def precision_score(y_true, y_pred, average='micro'): true_entities = set(y_true) pred_entities = set(y_pred) nb_correct = len(true_entities & pred_entities) nb_pred = len(pred_entities) score = nb_correct / nb_pred if nb_pred > 0 else 0 return score def recall_score(y_true, y_pred, average='micro', suffix=False): true_entities = set(y_true) pred_entities = set(y_pred) nb_correct = len(true_entities & pred_entities) nb_true = len(true_entities) score = nb_correct / nb_true if nb_true > 0 else 0 return score def f_score(y_true, y_pred, average='micro', suffix=False): true_entities = set(y_true) pred_entities = set(y_pred) nb_correct = len(true_entities & pred_entities) nb_pred = len(pred_entities) nb_true = len(true_entities) p = nb_correct / nb_pred if nb_pred > 0 else 0 r = nb_correct / nb_true if nb_true > 0 else 0 score = 2 * p * r / (p + r) if p + r > 0 else 0 return score def compute_metrics(p: EvalPrediction, mode, dateset=None) -> Dict: preds_list, out_label_list = align_predictions(p.predictions, p.label_ids, dateset) words_list, preds_list = post_align_predictions( dateset, preds_list, tokenizer) if mode == "dev": reader_file = os.path.join(data_args.data_dir, f"{mode}.txt") elif mode == "test": reader_file = os.path.join(data_args.test_data_dir, f"{mode}.txt") else: raise ValueError(f"mode is error: {mode}") if mode == "dev": data_dir = TRAIN_DIR elif mode == "test": data_dir = TEST_DIR else: raise ValueError(f"mode is error: {mode}") with open(reader_file, "r", encoding="utf8") as reader: pre_tuple, real_tuple = conver_entity_list_to_tuple( reader, preds_list, words_list, data_dir) # print(real_tuple[0], pre_tuple[0]) result = { # "accuracy_score": accuracy_score(out_label_list, preds_list), "all_precision": precision_score(real_tuple, pre_tuple), "all_recall": recall_score(real_tuple, pre_tuple), "all_f_score": f_score(real_tuple, pre_tuple), } for label in LABEL_LIST: sub_pre_tuple = [t for t in pre_tuple[:] if t[-1] == label] sub_real_tuple = [t for t in real_tuple[:] if t[-1] == label] result.update({ f"{label}_precision": precision_score(sub_real_tuple, sub_pre_tuple), f"{label}_recall": recall_score(sub_real_tuple, sub_pre_tuple), f"{label}_f_score": f_score(sub_real_tuple, sub_pre_tuple), }) metrics_report = f"\n{'Tag':20s}\t{'Precision':9s}\t{'Recall':9s}\t{'F-Score':9s}\t\n" for label in LABEL_LIST: pkey = f"{label}_precision" rkey = f"{label}_recall" fkey = f"{label}_f_score" p, r, f = result[pkey], result[rkey], result[fkey] metrics_report += f"{label:20s}\t{p:9.7f}\t{r:9.7f}\t{f:9.7f}\t\n" metrics_report += "<BLANKLINE>\n" pkey = f"all_precision" rkey = f"all_recall" fkey = f"all_f_score" p, r, f = result[pkey], result[rkey], result[fkey] label = "ALL" metrics_report += f"{label:20s}\t{p:9.7f}\t{r:9.7f}\t{f:9.7f}\t\n" logger.info("--------metricd report---------") logger.info(metrics_report) return result # Initialize our Trainer trainer = Trainer( model=model, args=training_args, train_dataset=train_dataset, eval_dataset=eval_dataset, compute_metrics=compute_metrics, ) # Training if training_args.do_train: trainer.train(model_path=model_args.model_name_or_path if os.path. isdir(model_args.model_name_or_path) else None) trainer.save_model() # For convenience, we also re-save the tokenizer to the same directory, # so that you can share your model easily on huggingface.co/models =) if trainer.is_world_master(): tokenizer.save_pretrained(training_args.output_dir) select_model(training_args.output_dir, k_fold=10, select=True, rm=False) if training_args.do_eval or training_args.do_predict or training_args.do_predict_dev: model = AutoModelForTokenClassification.from_pretrained( training_args.output_dir, from_tf=bool(".ckpt" in model_args.model_name_or_path), config=config, cache_dir=model_args.cache_dir, training_args=training_args) trainer = Trainer( model=model, args=training_args, train_dataset=train_dataset, eval_dataset=eval_dataset, compute_metrics=compute_metrics, ) logger.warning(f"load best mode from {training_args.output_dir}") # Evaluation results = {} if training_args.do_eval: logger.info("*** Evaluate ***") result = trainer.evaluate() output_eval_file = os.path.join(training_args.output_dir, "eval_results.txt") if trainer.is_world_master(): with open(output_eval_file, "w") as writer: logger.info("***** Eval results *****") for key, value in result.items(): # logger.info(" %s = %s", key, value) writer.write("%s = %s\n" % (key, value)) results.update(result) # Predict if training_args.do_predict: test_dataset = TokenClassificationDataset( token_classification_task=token_classification_task, data_dir=data_args.test_data_dir, tokenizer=tokenizer, labels=labels, model_type=config.model_type, max_seq_length=data_args.max_seq_length, overwrite_cache=data_args.overwrite_cache, mode=Split.test, training_args=training_args, remove_label=Remove_Label) # logger.warning(list(zip(tokens_list[3], test_dataset.features[3].input_ids))) predictions, label_ids, metrics = trainer.predict(test_dataset) preds_list, _ = align_predictions(predictions, label_ids, test_dataset) # 数据对齐 words_list, preds_list = post_align_predictions( test_dataset, preds_list, tokenizer) output_test_results_file = os.path.join(training_args.output_dir, "test_results.txt") if trainer.is_world_master(): with open(output_test_results_file, "w") as writer: for key, value in metrics.items(): logger.info(" %s = %s", key, value) writer.write("%s = %s\n" % (key, value)) # Save predictions output_test_predictions_file = os.path.join(training_args.output_dir, "test_predictions.txt") if trainer.is_world_master(): with open(output_test_predictions_file, "w") as writer: with open(os.path.join(data_args.test_data_dir, "test.txt"), "r") as f: token_classification_task.write_predictions_to_file( writer, f, preds_list, words_list) # Predict dev if training_args.do_predict_dev: # logger.warning(list(zip(tokens_list[3], test_dataset.features[3].input_ids))) predictions, label_ids, metrics = trainer.predict( eval_dataset, description="Evaluation") preds_list, _ = align_predictions(predictions, label_ids, eval_dataset) # 数据对齐 words_list, preds_list = post_align_predictions( eval_dataset, preds_list, tokenizer) output_test_results_file = os.path.join(training_args.output_dir, "dev_results.txt") if trainer.is_world_master(): with open(output_test_results_file, "w") as writer: for key, value in metrics.items(): writer.write("%s = %s\n" % (key, value)) # Save predictions output_test_predictions_file = os.path.join(training_args.output_dir, "dev_predictions.txt") if trainer.is_world_master(): with open(output_test_predictions_file, "w") as writer: with open(os.path.join(data_args.data_dir, "dev.txt"), "r") as f: token_classification_task.write_predictions_to_file( writer, f, preds_list, words_list) return results