diff --git a/xcore/data/datasets.py b/xcore/data/datasets.py index 442f0b5..0535260 100644 --- a/xcore/data/datasets.py +++ b/xcore/data/datasets.py @@ -21,7 +21,7 @@ def __init__(self, name: str, path: str, batch_size, processed_dataset_path, tok self.stage = name self.path = path self.batch_size = batch_size - self.tokenizer = AutoTokenizer.from_pretrained(tokenizer, use_fast=True, add_prefix_space=True) + self.tokenizer = AutoTokenizer.from_pretrained(tokenizer, use_fast=False, add_prefix_space=True) self.max_doc_len = kwargs.get("max_doc_len", None) self.cross = kwargs.get("type", None) == "cross" self.book = kwargs.get("type", None) == "book" diff --git a/xcore/evaluate.py b/xcore/evaluate.py index 20a6960..5c512f2 100644 --- a/xcore/evaluate.py +++ b/xcore/evaluate.py @@ -1,80 +1,159 @@ +import os +import re import json -import hydra +import operator import subprocess +import collections +from pprint import pprint +import hydra import torch -from omegaconf import omegaconf -from xcore.common.util import * -from xcore.common.metrics import * from tqdm import tqdm -from data.pl_data_modules import CrossDataModule +from omegaconf import omegaconf from models.pl_modules import CrossPLModule +from data.pl_data_modules import CrossDataModule +from xcore.common.util import (extract_mentions_to_clusters, + original_token_offsets3) +from xcore.common.metrics import OfficialCoNLL2012CorefEvaluator from xcore.utils.loggingl import get_console_logger logger = get_console_logger() +# extract doc_key and part given format 'doc_key.p.0' +DOC_KEY_RE = re.compile(r'^(.+)\.p\.(\d+)$') def jsonlines_to_html(jsonlines_input_name, output): cwd = str(hydra.utils.get_original_cwd()) - subprocess.call( - "python3 " - + cwd - + "/xcore/utils/corefconversion/jsonlines2text.py " - + cwd - + "/" - + jsonlines_input_name - + " -i -o " - + cwd - + "/experiments/" - + output - + ".html --sing-color" - ' "black" --cm "common"', - shell=True, - ) + subprocess.check_call( + ["python3", + cwd + "/xcore/utils/corefconversion/jsonlines2text.py", + jsonlines_input_name, + "-i", + "-o", output, + "--sing-color", "black", + "--cm", "common"], + shell=False) + + +def output_conll(info, output_file): + # Based on https://github.com/kentonl/e2e-coref/blob/master/conll.py + with open(output_file, 'w', encoding='utf8') as out: + for infos in info: + doc_key = infos['doc_key'] + clusters = infos['clusters'] + start_map = collections.defaultdict(list) + end_map = collections.defaultdict(list) + word_map = collections.defaultdict(list) + for cluster_id, mentions in enumerate(clusters): + for start, end in mentions: + if start == end: + word_map[start].append(cluster_id) + else: + start_map[start].append((cluster_id, end)) + end_map[end].append((cluster_id, start)) + for k,v in start_map.items(): + start_map[k] = [cluster_id for cluster_id, end + in sorted(v, key=operator.itemgetter(1), reverse=True)] + for k,v in end_map.items(): + end_map[k] = [cluster_id for cluster_id, start + in sorted(v, key=operator.itemgetter(1), reverse=True)] + + docname, part = DOC_KEY_RE.match(doc_key).groups() + print(f'#begin document ({docname}); part {part}', file=out) + word_index = 0 + for sentence in infos['sentences']: + for tokenid, token in enumerate(sentence): + row = [docname, part, tokenid, token] + ['-'] * 8 + coref_list = [] + if word_index in end_map: + for cluster_id in end_map[word_index]: + coref_list.append(f'{cluster_id})') + if word_index in word_map: + for cluster_id in word_map[word_index]: + coref_list.append(f'({cluster_id})') + if word_index in start_map: + for cluster_id in start_map[word_index]: + coref_list.append(f'({cluster_id}') + row[-1] = '|'.join(coref_list) if coref_list else '-' + print(*row, sep='\t', file=out) + word_index += 1 + print('', file=out) + print('#end document', file=out) @torch.no_grad() def evaluate(conf: omegaconf.DictConfig): device = conf.evaluation.device - - hydra.utils.log.info("Using {} as device".format(device)) - pl_data_module: CrossDataModule = hydra.utils.instantiate(conf.data.datamodule, _recursive_=False) + + hydra.utils.log.info('Using %s as device', device) + pl_data_module: CrossDataModule = hydra.utils.instantiate( + conf.data.datamodule, _recursive_=False) pl_data_module.prepare_data() pl_data_module.setup("test") - jsonlines_to_html(pl_data_module.test_dataloader().dataset.path, "test") - + cwd = str(hydra.utils.get_original_cwd()) + # 'data/mydataset/dev.jsonlines' -> 'mydataset' + dataset = os.path.basename(os.path.dirname( + pl_data_module.test_dataloader().dataset.path)) + # Write all output to a sibling directory of the checkpoints directory + # experiments/xcore/myexperiment/wandb/run-2026{...}/files/{dataset} + # This means we can evaluate a single checkpoint on multiple test datasets + # and keep the results in separate directories. + outpath = os.path.dirname(conf.evaluation.checkpoint + ).removesuffix('checkpoints') + dataset + if not os.path.exists(outpath): + os.mkdir(outpath) + # All output filenames are of the form {subset}_{modality}, + # where modality is gold or output. + jsonlines_to_html( + cwd + '/' + pl_data_module.test_dataloader().dataset.path, + outpath + "/test_gold.html") logger.log(f"Instantiating the Model from {conf.evaluation.checkpoint}") - model = CrossPLModule.load_from_checkpoint(conf.evaluation.checkpoint, _recursive_=False, map_location=device) + model = CrossPLModule.load_from_checkpoint( + conf.evaluation.checkpoint, _recursive_=False, + map_location=device, weights_only=False) + gold = [] info = [] - with open(hydra.utils.get_original_cwd() + "/" + pl_data_module.test_dataloader().dataset.path, "r") as f: - for line in f.readlines(): + with open(cwd + "/" + pl_data_module.test_dataloader().dataset.path, + 'r', encoding='utf8') as infile: + for line in infile.readlines(): doc = json.loads(line) if "sentences" in doc: - info.append({"doc_key": doc["doc_key"], "sentences": doc["sentences"]}) + info.append({ + "doc_key": doc["doc_key"], + "sentences": doc["sentences"]}) clusters = [] if "clusters" in doc: for cluster in doc["clusters"]: if not conf.evaluation.singletons and len(cluster) < 2: continue - clusters.append(tuple([(m[0], m[1]) for m in cluster])) + clusters.append(tuple((m[0], m[1]) for m in cluster)) gold.append(clusters) - mention_to_gold_clusters = [extract_mentions_to_clusters([tuple(g) for g in gold_element]) for gold_element in gold] + mention_to_gold_clusters = [ + extract_mentions_to_clusters([tuple(g) for g in gold_element]) + for gold_element in gold] + + predictions = model_predictions_with_dataloader( + model, pl_data_module.test_dataloader(), device, + conf.evaluation.singletons) + mention_to_predicted_clusters = [extract_mentions_to_clusters(p) + for p in predictions] - device_and_singletons ={"device": device, - "singletons": conf.evaluation.singletons} - predictions = model_predictions_with_dataloader(model, pl_data_module.test_dataloader(), device_and_singletons) - mention_to_predicted_clusters = [extract_mentions_to_clusters(p) for p in predictions] + pprint(evaluate_coref_scores(predictions, gold, + mention_to_predicted_clusters, mention_to_gold_clusters), + sort_dicts=False) - print(evaluate_coref_scores(predictions, gold, mention_to_predicted_clusters, mention_to_gold_clusters)) + with open(outpath + '/test_output.jsonlines', 'w', + encoding='utf8') as outfile: + for pred, infos in zip(predictions, info): + infos["clusters"] = pred + outfile.write(json.dumps(infos) + "\n") - with open(hydra.utils.get_original_cwd() + "/experiments/output.jsonlines", "w") as f: - for pred, infos in zip(predictions, info): - infos["clusters"] = pred - f.write(json.dumps(infos) + "\n") + jsonlines_to_html( + outpath + "/test_output.jsonlines", + outpath + "/test_output.html") - jsonlines_to_html("experiments/output.jsonlines", "output") - return + output_conll(info, outpath + '/test_output.conll') def evaluate_coref_scores(pred, gold, mention_to_pred, mention_to_gold): @@ -82,39 +161,44 @@ def evaluate_coref_scores(pred, gold, mention_to_pred, mention_to_gold): for p, g, m2p, m2g in zip(pred, gold, mention_to_pred, mention_to_gold): evaluator.update(p, g, m2p, m2g) - result = [] + result = {} for metric in ["muc", "b_cubed", "ceafe", "conll2012"]: - result.append(dict(zip(["precision", "recall", "f1_score"], evaluator.get_prf(metric)))) + result[metric] = dict(zip([ + "precision", "recall", "f1_score"], + evaluator.get_prf(metric))) return result -def model_predictions_with_dataloader(model, test_dataloader, device_and_singletons): - model.to(device_and_singletons["device"]) +def model_predictions_with_dataloader( + model, test_dataloader, device, singletons): + model.to(device) model.eval() predictions = [] - for batch in tqdm(test_dataloader, desc="Test", total=test_dataloader.__len__()): + for batch in tqdm( + test_dataloader, desc="Test", total=len(test_dataloader)): output = model.model( stage="temp", - input_ids=[elem.to(device_and_singletons["device"]) for elem in batch["index_input_ids"]], - attention_mask=[elem.to(device_and_singletons["device"]) for elem in batch["index_attention_mask"]], - eos_mask=[elem.to(device_and_singletons["device"]) for elem in batch["index_eos_mask"]], - gold_starts=[elem.to(device_and_singletons["device"]) for elem in batch["index_gold_starts"]], - gold_mentions=[elem.to(device_and_singletons["device"]) for elem in batch["index_gold_mentions"]], + input_ids=[elem.to(device) for elem in batch["index_input_ids"]], + attention_mask=[elem.to(device) + for elem in batch["index_attention_mask"]], + eos_mask=[elem.to(device) for elem in batch["index_eos_mask"]], + gold_starts=[elem.to(device) + for elem in batch["index_gold_starts"]], + gold_mentions=[elem.to(device) + for elem in batch["index_gold_mentions"]], gold_clusters=batch["index_gold_clusters"], - singletons=device_and_singletons["singletons"], - full_clusters=batch["gold_c"].to(device_and_singletons["device"]), + singletons=singletons, + full_clusters=batch["gold_c"].to(device), temp=batch["temp"], tokens=batch["t_tokens"], subtoken_map=batch["t_subtoken_map"], - new_token_map=batch["t_new_token_map"], - ) + new_token_map=batch["t_new_token_map"]) clusters_predicted = original_token_offsets3( clusters=output["pred_dict"]["full_coreferences"], subtoken_map=batch["subtoken_map"][0], - new_token_map=batch["new_token_map"][0], - ) + new_token_map=batch["new_token_map"][0]) predictions.append(clusters_predicted) return predictions @@ -125,5 +209,5 @@ def main(conf: omegaconf.DictConfig): evaluate(conf) -if __name__ == "__main__": +if __name__ == '__main__': main() diff --git a/xcore/models/model_cross.py b/xcore/models/model_cross.py index ff87db7..9b435bd 100644 --- a/xcore/models/model_cross.py +++ b/xcore/models/model_cross.py @@ -20,12 +20,12 @@ ) class attention(torch.nn.Module): - def __init__(self, model, representation): + def __init__(self, model, representation, input_dim): super().__init__() self.model = model self.t = RepresentationLayer( type="FC", # fullyconnected - input_dim=2048, + input_dim=input_dim, output_dim=768, hidden_dim=1024, ) @@ -45,9 +45,13 @@ def __init__(self, *args, **kwargs): super().__init__() # document transformer encoder self.encoder_hf_model_name = kwargs["huggingface_model_name"] - self.encoder = AutoModel.from_pretrained(self.encoder_hf_model_name) + self.encoder = AutoModel.from_pretrained(self.encoder_hf_model_name).train() self.encoder_config = AutoConfig.from_pretrained(self.encoder_hf_model_name) - self.encoder.resize_token_embeddings(self.encoder.embeddings.word_embeddings.num_embeddings + 3) + try: + num_embeddings = self.encoder.embeddings.word_embeddings.num_embeddings + except AttributeError: + num_embeddings = self.encoder.embeddings.tok_embeddings.num_embeddings + self.encoder.resize_token_embeddings(num_embeddings + 3) self.device = self.encoder.device # freeze @@ -187,7 +191,7 @@ def __init__(self, *args, **kwargs): self.cluster_model = DistilBertModel(self.cluster_model_config).to(self.encoder.device) self.cluster_model.transformer.layer = self.cluster_model.transformer.layer[: self.cluster_model_num_layers] self.cluster_model.embeddings.word_embeddings = None - self.cluster_transformer = attention(model=self.cluster_model, representation=self.cluster_representation) + self.cluster_transformer = attention(model=self.cluster_model, representation=self.cluster_representation, input_dim=self.mention_hidden_size) self.antecedent_coref_classifier = RepresentationLayer( type=self.representation_layer_type, # fullyconnected diff --git a/xcore/models/xcore_model.py b/xcore/models/xcore_model.py index 004ff65..d90a9c1 100644 --- a/xcore/models/xcore_model.py +++ b/xcore/models/xcore_model.py @@ -36,7 +36,7 @@ def __get_model_path__(self, hf_name_or_path): return path def __get_model_tokenizer__(self): - tokenizer = AutoTokenizer.from_pretrained(self.model.encoder_hf_model_name, use_fast=True, add_prefix_space=True) + tokenizer = AutoTokenizer.from_pretrained(self.model.encoder_hf_model_name, use_fast=False, add_prefix_space=True) special_tokens_dict = {"additional_special_tokens": ["[SPEAKER_START]", "[SPEAKER_END]"]} tokenizer.add_special_tokens(special_tokens_dict) return tokenizer diff --git a/xcore/train.py b/xcore/train.py index 681dfcf..c92fea7 100644 --- a/xcore/train.py +++ b/xcore/train.py @@ -1,4 +1,5 @@ import os +import warnings import hydra import omegaconf import pytorch_lightning as pl @@ -24,6 +25,8 @@ def train(conf: omegaconf.DictConfig) -> None: + # don't warn about multiprocessing for data loaders + warnings.filterwarnings("ignore", ".*does not have many workers.*") # fancy logger console = Console() # reproducibility @@ -46,6 +49,14 @@ def train(conf: omegaconf.DictConfig) -> None: # conf.train.pl_trainer.accelerator = "cpu" + # The following was prompted by this log message: + # You are using a CUDA device ('NVIDIA A100 80GB PCIe') that has Tensor + # Cores. To properly utilize them, you should set + # `torch.set_float32_matmul_precision('medium' | 'high')` which will + # trade-off precision for performance. For more details, read + # https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision + torch.set_float32_matmul_precision('medium') + # data module declaration console.log(f"Instantiating the Data Module") @@ -98,7 +109,8 @@ def train(conf: omegaconf.DictConfig) -> None: trainer.fit(pl_module, datamodule=pl_data_module) # module test - trainer.test(pl_module, datamodule=pl_data_module) + # disabled due to bugs + # trainer.test(pl_module, datamodule=pl_data_module) def set_determinism_the_old_way(deterministic: bool):