From e1aa3d32dd5de0cc8b163a3e5ddeaaf47813cadc Mon Sep 17 00:00:00 2001 From: andreasvc Date: Mon, 20 Apr 2026 14:35:46 +0200 Subject: [PATCH 1/4] Disable fast tokenizer to avoid warnings --- xcore/data/datasets.py | 2 +- xcore/models/xcore_model.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) 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/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 From 69a34ae39c8030a192704f949e20a86d3a208584 Mon Sep 17 00:00:00 2001 From: andreasvc Date: Mon, 20 Apr 2026 14:58:19 +0200 Subject: [PATCH 2/4] Support embeddings with different dimensions - The attention() class had a hard-coded dimension of 2048 for the input, which works for Deberta, but not other models such as mmbert; the input dimension is now specified with a parameter. - Set the encoder model to training mode after loading https://huggingface.co/docs/transformers/model_doc/auto#transformers.AutoModelForPreTraining.from_pretrained - mmbert stores token embeddings in "self.encoder.tok_embeddings" instead of "self.encoder.word_embeddings"; both are now supported --- xcore/models/model_cross.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) 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 From cfaf1e7703095565c59652aced0a17cc671ec1fd Mon Sep 17 00:00:00 2001 From: andreasvc Date: Mon, 20 Apr 2026 15:06:41 +0200 Subject: [PATCH 3/4] Suppress warnings, disable test mode - Suppress warnings about multiprocessing for data loaders; loading the data is not a bottleneck so the warnings are unnecessary - Depending on the GPU, you may get warnings about setting matmul_precision. Added code to set matmul precision to medium, which seems like a good trade-off (but your mileage may vary on different hardware) - Disable evaluation on the test set during training. Apparently, the code has bugs and is not expected to work. Therefore it is now disabled, to avoid giving the impression that there is an actual issue with training the model. Evaluation is supposed to be performed with the evaluate.py script. --- xcore/train.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) 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): From 0a6df44dc7276bc49ecca0bf778962ae937e76a7 Mon Sep 17 00:00:00 2001 From: andreasvc Date: Mon, 20 Apr 2026 15:16:24 +0200 Subject: [PATCH 4/4] Improved evaluate.py script - store the .conll output of the model, useful if you want to run other evaluation scripts on the output. - write output to a separate directory and use filenames of the form '{subset}_{modality}', e.g. 'test_output' and 'test_gold' to clearly indicate the type of file. The output is written to a directory based on the dataset: 'experiments/xcore/myexperiment/wandb/run-2026{...}/files/{dataset}' A model can therefore be evaluated on multiple datatsets. - pretty-print evaluation results --- xcore/evaluate.py | 202 ++++++++++++++++++++++++++++++++-------------- 1 file changed, 143 insertions(+), 59 deletions(-) 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()