def recognize(test_file_name: str, split_by_paragraphs: bool, recognizer: BERT_NER, results_file_name: str): X_test, y_test = load_dataset_from_bio( test_file_name, paragraph_separators=({'-DOCSTART-'} if split_by_paragraphs else None), stopwords={'-DOCSTART-'}) print('The CoNLL-2003 data for final testing have been loaded...') print('Number of samples is {0}.'.format(len(y_test))) print('') y_pred = recognizer.predict(X_test) f1, precision, recall, quality_by_entities = calculate_prediction_quality( y_test, y_pred, classes_list=recognizer.classes_list_) print('All entities:') print(' F1-score is {0:.2%}.'.format(f1)) print(' Precision is {0:.2%}.'.format(precision)) print(' Recall is {0:.2%}.'.format(recall)) for ne_type in sorted(list(quality_by_entities.keys())): print(' {0}'.format(ne_type)) print(' F1-score is {0:.2%}.'.format( quality_by_entities[ne_type][0])) print(' Precision is {0:.2%}.'.format( quality_by_entities[ne_type][1])) print(' Recall is {0:.2%}.'.format(quality_by_entities[ne_type][2])) print('') save_dataset_as_bio(test_file_name, X_test, y_pred, results_file_name, stopwords={'-DOCSTART-'})
def train(train_file_name: str, valid_file_name: str, split_by_paragraphs: bool, bert_will_be_tuned: bool, lstm_layer_size: Union[int, None], l2: float, max_epochs: int, batch_size: int, gpu_memory_frac: float, model_name: str) -> BERT_NER: if os.path.isfile(model_name): with open(model_name, 'rb') as fp: recognizer = pickle.load(fp) assert isinstance(recognizer, BERT_NER) print('The NER has been successfully loaded from the file `{0}`...'. format(model_name)) print('') else: X_train, y_train = load_dataset_from_bio( train_file_name, paragraph_separators=({'-DOCSTART-'} if split_by_paragraphs else None), stopwords={'-DOCSTART-'}) X_val, y_val = load_dataset_from_bio( valid_file_name, paragraph_separators=({'-DOCSTART-'} if split_by_paragraphs else None), stopwords={'-DOCSTART-'}) print( 'The CoNLL-2003 data for training and validation have been loaded...' ) print('Number of samples for training is {0}.'.format(len(y_train))) print('Number of samples for validation is {0}.'.format(len(y_val))) print('') if BERT_NER.PATH_TO_BERT is None: bert_hub_module_handle = 'https://tfhub.dev/google/bert_cased_L-12_H-768_A-12/1' else: bert_hub_module_handle = None recognizer = BERT_NER(finetune_bert=bert_will_be_tuned, batch_size=batch_size, l2_reg=l2, bert_hub_module_handle=bert_hub_module_handle, lstm_units=lstm_layer_size, max_epochs=max_epochs, patience=5, gpu_memory_frac=gpu_memory_frac, verbose=True, random_seed=42, lr=1e-6 if bert_will_be_tuned else 1e-4) recognizer.fit(X_train, y_train, validation_data=(X_val, y_val)) print('') print( 'The NER has been successfully fitted and saved into the file `{0}`...' .format(model_name)) y_pred = recognizer.predict(X_val) f1, precision, recall, quality_by_entities = calculate_prediction_quality( y_val, y_pred, classes_list=recognizer.classes_list_) print('All entities:') print(' F1-score is {0:.2%}.'.format(f1)) print(' Precision is {0:.2%}.'.format(precision)) print(' Recall is {0:.2%}.'.format(recall)) for ne_type in sorted(list(quality_by_entities.keys())): print(' {0}'.format(ne_type)) print(' F1-score is {0:.2%}.'.format( quality_by_entities[ne_type][0])) print(' Precision is {0:.2%}.'.format( quality_by_entities[ne_type][1])) print(' Recall is {0:.2%}.'.format( quality_by_entities[ne_type][2])) print('') with open(model_name, 'wb') as fp: pickle.dump(recognizer, fp) return recognizer