一个极速且轻量的信息抽取模型,用于实体链接和关系抽取。
🛠️ 安装
从 PyPI 安装
pip install relik
其他安装选项
安装可选依赖项
安装所有可选依赖项。
pip install relik[all]
安装用于训练和评估的可选依赖项。
pip install relik[train]
使用可选依赖项安装 FAISS
FAISS PyPI 包仅支持 CPU。对于 GPU,请从源代码安装或使用 conda 包。
对于 CPU:
pip install relik[faiss]
对于 GPU:
conda create -n relik python=3.10
conda activate relik
# install pytorch
conda install -y pytorch=2.1.0 pytorch-cuda=12.1 -c pytorch -c nvidia
# GPU
conda install -y -c pytorch -c nvidia faiss-gpu=1.8.0
# or GPU with NVIDIA RAFT
conda install -y -c pytorch -c nvidia -c rapidsai -c conda-forge faiss-gpu-raft=1.8.0
pip install relik
使用可选依赖项安装,以便使用 FastAPI 和 Ray 提供模型服务。
pip install relik[serve]
从源码安装
git clone https://github.com/SapienzaNLP/relik.git
cd relik
pip install -e .[all]
🤖 模型
-
ReLiK Large for Relation Extraction (🆕 Large RE v2, Colab ✅):
relik-ie/relik-relation-extraction-large -
ReLiK Large for Closed Information Extraction (🆕 large EL + RE, Colab ✅):
https://huggingface.co/relik-ie/relik-cie-large -
ReLiK Extra Large for Closed Information Extraction (🆕 Our thicc boi for EL + RE):
relik-ie/relik-cie-xl -
ReLiK Small for Entity Linking (🆕🤏⚡ Tiny and Fast EL, Colab ✅):
sapienzanlp/relik-entity-linking-small -
ReLiK Small for Entity Linking (⚡ Small and Fast EL):
sapienzanlp/relik-entity-linking-small -
ReLiK Small for Closed Information Extraction (🔥 EL + RE):
relik-ie/relik-cie-small -
ReLiK Large for Entity Linking (🔥 EL for the wild):
relik-ie/relik-entity-linking-large-robust -
ReLiK Small for Entity Linking (🔥 RE + NER):
relik-ie/relik-relation-extraction-small-wikipedia-ner
论文中的模型:
- ReLiK Large for Entity Linking (📝 Paper version):
sapienzanlp/relik-entity-linking-large - ReLik Base for Entity Linking (📝 Paper version):
sapienzanlp/relik-entity-linking-base - ReLiK Large for Relation Extraction (📝 Paper version):
sapienzanlp/relik-relation-extraction-nyt-large
完整的模型列表可在 🤗 Hugging Face 上找到。
其他模型尺寸将在未来提供 👀。
🚀 快速开始
ReLiK 是一个用于实体链接和关系抽取的轻量级且快速的模型。
它由两个主要组件组成:一个检索器和一个阅读器。
检索器负责从大型文档集合中检索相关文档,
而阅读器负责从检索到的文档中提取实体和关系。
ReLiK 可以使用 from_pretrained 方法来加载预训练管道。
以下是如何使用 ReLiK 进行实体链接的示例:
from relik import Relik
from relik.inference.data.objects import RelikOutput
relik = Relik.from_pretrained("sapienzanlp/relik-entity-linking-large")
relik_out: RelikOutput = relik("Michael Jordan was one of the best players in the NBA.")
RelikOutput( text="Michael Jordan was one of the best players in the NBA.", tokens=['Michael', 'Jordan', 'was', 'one', 'of', 'the', 'best', 'players', 'in', 'the', 'NBA', '.'], id=0, spans=[ Span(start=0, end=14, label="Michael Jordan", text="Michael Jordan"), Span(start=50, end=53, label="National Basketball Association", text="NBA"), ], triples=[], candidates=Candidates( span=[ [ [ {"text": "Michael Jordan", "id": 4484083}, {"text": "National Basketball Association", "id": 5209815}, {"text": "Walter Jordan", "id": 2340190}, {"text": "Jordan", "id": 3486773}, {"text": "50 Greatest Players in NBA History", "id": 1742909}, ... ] ] ] ), )
以及用于关系抽取:
from relik import Relik
from relik.inference.data.objects import RelikOutput
relik = Relik.from_pretrained("sapienzanlp/relik-relation-extraction-nyt-large")
relik_out: RelikOutput = relik("Michael Jordan was one of the best players in the NBA.")
RelikOutput( text='Michael Jordan was one of the best players in the NBA.', tokens=Michael Jordan was one of the best players in the NBA., id=0, spans=[ Span(start=0, end=14, label='--NME--', text='Michael Jordan'), Span(start=50, end=53, label='--NME--', text='NBA') ], triplets=[ Triplets( subject=Span(start=0, end=14, label='--NME--', text='Michael Jordan'), label='company', object=Span(start=50, end=53, label='--NME--', text='NBA'), confidence=1.0 ) ], candidates=Candidates( span=[], triplet=[ [ [ {"text": "company", "id": 4, "metadata": {"definition": "company of this person"}}, {"text": "nationality", "id": 10, "metadata": {"definition": "nationality of this person or entity"}}, {"text": "child", "id": 17, "metadata": {"definition": "child of this person"}}, {"text": "founded by", "id": 0, "metadata": {"definition": "founder or co-founder of this organization, religion or place"}}, {"text": "residence", "id": 18, "metadata": {"definition": "place where this person has lived"}}, ... ] ] ] ), )
用法
检索器和阅读器可以单独使用。 在仅使用检索器的 ReLiK 情况下,输出将包含输入文本的候选项。
仅检索器示例:
from relik import Relik
from relik.inference.data.objects import RelikOutput
# If you want to use only the retriever
retriever = Relik.from_pretrained("sapienzanlp/relik-entity-linking-large", reader=None)
relik_out: RelikOutput = retriever("Michael Jordan was one of the best players in the NBA.")
RelikOutput( text="Michael Jordan was one of the best players in the NBA.", tokens=['Michael', 'Jordan', 'was', 'one', 'of', 'the', 'best', 'players', 'in', 'the', 'NBA', '.'], id=0, spans=[], triples=[], candidates=Candidates( span=[ [ {"text": "Michael Jordan", "id": 4484083}, {"text": "National Basketball Association", "id": 5209815}, {"text": "Walter Jordan", "id": 2340190}, {"text": "Jordan", "id": 3486773}, {"text": "50 Greatest Players in NBA History", "id": 1742909}, ... ] ], triplet=[], ), )
仅读者示例:
from relik import Relik
from relik.inference.data.objects import RelikOutput
# If you want to use only the reader
reader = Relik.from_pretrained("sapienzanlp/relik-entity-linking-large", retriever=None)
candidates = [
"Michael Jordan",
"National Basketball Association",
"Walter Jordan",
"Jordan",
"50 Greatest Players in NBA History",
]
text = "Michael Jordan was one of the best players in the NBA."
relik_out: RelikOutput = reader(text, candidates=candidates)
RelikOutput( text="Michael Jordan was one of the best players in the NBA.", tokens=['Michael', 'Jordan', 'was', 'one', 'of', 'the', 'best', 'players', 'in', 'the', 'NBA', '.'], id=0, spans=[ Span(start=0, end=14, label="Michael Jordan", text="Michael Jordan"), Span(start=50, end=53, label="National Basketball Association", text="NBA"), ], triples=[], candidates=Candidates( span=[ [ [ { "text": "Michael Jordan", "id": -731245042436891448, }, { "text": "National Basketball Association", "id": 8135443493867772328, }, { "text": "Walter Jordan", "id": -5873847607270755146, "metadata": {}, }, {"text": "Jordan", "id": 6387058293887192208, "metadata": {}}, { "text": "50 Greatest Players in NBA History", "id": 2173802663468652889, }, ] ] ], ), )
CLI
ReLiK 提供了一个 CLI,用于为模型提供 FastAPI 服务器,或对数据集执行推理。
relik serve
relik serve --help
Usage: relik serve [OPTIONS] RELIK_PRETRAINED [DEVICE] [RETRIEVER_DEVICE]
[DOCUMENT_INDEX_DEVICE] [READER_DEVICE] [PRECISION]
[RETRIEVER_PRECISION] [DOCUMENT_INDEX_PRECISION]
[READER_PRECISION] [ANNOTATION_TYPE]
╭─ Arguments ─────────────────────────────────────────────────────────────────────────────────────────╮
│ * relik_pretrained TEXT [default: None] [required] │
│ device [DEVICE] The device to use for relik (e.g., │
│ 'cuda', 'cpu'). │
│ [default: None] │
│ retriever_device [RETRIEVER_DEVICE] The device to use for the retriever │
│ (e.g., 'cuda', 'cpu'). │
│ [default: None] │
│ document_index_device [DOCUMENT_INDEX_DEVICE] The device to use for the index │
│ (e.g., 'cuda', 'cpu'). │
│ [default: None] │
│ reader_device [READER_DEVICE] The device to use for the reader │
│ (e.g., 'cuda', 'cpu'). │
│ [default: None] │
│ precision [PRECISION] The precision to use for relik │
│ (e.g., '32', '16'). │
│ [default: 32] │
│ retriever_precision [RETRIEVER_PRECISION] The precision to use for the │
│ retriever (e.g., '32', '16'). │
│ [default: None] │
│ document_index_precision [DOCUMENT_INDEX_PRECISION] The precision to use for the index │
│ (e.g., '32', '16'). │
│ [default: None] │
│ reader_precision [READER_PRECISION] The precision to use for the reader │
│ (e.g., '32', '16'). │
│ [default: None] │
│ annotation_type [ANNOTATION_TYPE] The type of annotation to use (e.g., │
│ 'CHAR', 'WORD'). │
│ [default: char] │
╰─────────────────────────────────────────────────────────────────────────────────────────────────────╯
╭─ Options ───────────────────────────────────────────────────────────────────────────────────────────╮
│ --host TEXT [default: 0.0.0.0] │
│ --port INTEGER [default: 8000] │
│ --frontend --no-frontend [default: no-frontend] │
│ --help Show this message and exit. │
╰─────────────────────────────────────────────────────────────────────────────────────────────────────╯
例如:
relik serve sapienzanlp/relik-entity-linking-large
relik inference
relik inference --help
Usage: relik inference [OPTIONS] MODEL_NAME_OR_PATH INPUT_PATH OUTPUT_PATH
╭─ Arguments ─────────────────────────────────────────────────────────────────────────────────────────────╮
│ * model_name_or_path TEXT [default: None] [required] │
│ * input_path TEXT [default: None] [required] │
│ * output_path TEXT [default: None] [required] │
╰─────────────────────────────────────────────────────────────────────────────────────────────────────────╯
╭─ Options ───────────────────────────────────────────────────────────────────────────────────────────────╮
│ --batch-size INTEGER [default: 8] │
│ --num-workers INTEGER [default: 4] │
│ --device TEXT [default: cuda] │
│ --precision TEXT [default: fp16] │
│ --top-k INTEGER [default: 100] │
│ --window-size INTEGER [default: None] │
│ --window-stride INTEGER [default: None] │
│ --annotation-type TEXT [default: char] │
│ --progress-bar --no-progress-bar [default: progress-bar] │
│ --model-kwargs TEXT [default: None] │
│ --inference-kwargs TEXT [default: None] │
│ --help Show this message and exit. │
╰─────────────────────────────────────────────────────────────────────────────────────────────────────────╯
例如:
relik inference sapienzanlp/relik-entity-linking-large data.txt output.jsonl
Docker 镜像
ReLiK 的 Docker 镜像可在 Docker Hub 上获取。您可以使用以下命令拉取最新镜像:
docker pull sapienzanlp/relik:latest
并使用以下命令运行该镜像:
docker run -p 12345:8000 sapienzanlp/relik:latest -c relik-ie/relik-cie-small
该 API 将在 http://localhost:12345 处可用。它暴露了一个端点 /relik,该端点具有多个可传递给模型的参数。
API 的简要文档可在 http://localhost:12345/docs 处找到。以下是查询 API 的简单示例:
curl -X 'GET' \
'http://127.0.0.1:12345/api/relik?text=Michael%20Jordan%20was%20one%20of%20the%20best%20players%20in%20the%20NBA.&is_split_into_words=false&retriever_batch_size=32&reader_batch_size=32&return_windows=false&use_doc_topic=false&annotation_type=char&relation_threshold=0.5' \
-H 'accept: application/json'
以下是可以传递给 docker 镜像的完整参数列表:
docker run sapienzanlp/relik:latest -h
Usage: relik [-h --help] [-c --config] [-p --precision] [-d --device] [--retriever] [--retriever-device]
[--retriever-precision] [--index-device] [--index-precision] [--reader] [--reader-device] [--reader-precision]
[--annotation-type] [--frontend] [--workers] -- start the FastAPI server for the RElik model
where:
-h --help Show this help text
-c --config Pretrained ReLiK config name (from HuggingFace) or path
-p --precision Precision, default '32'.
-d --device Device to use, default 'cpu'.
--retriever Override retriever model name.
--retriever-device Override retriever device.
--retriever-precision Override retriever precision.
--index-device Override index device.
--index-precision Override index precision.
--reader Override reader model name.
--reader-device Override reader device.
--reader-precision Override reader precision.
--annotation-type Annotation type ('char', 'word'), default 'char'.
--frontend Whether to start the frontend server.
--workers Number of workers to use.
📚 开始之前
在以下章节中,我们提供了一份逐步指南,介绍如何准备数据、训练检索器和阅读器,以及评估模型。
实体链接
所有数据应具有以下结构:
{
"doc_id": int, # Unique identifier for the document
"doc_text": txt, # Text of the document
"doc_span_annotations": # Char level annotations
[
[start, end, label],
[start, end, label],
...
]
}
我们使用 BLINK (Wu et al., 2019) 和 AIDA (Hoffart et al, 2011) 数据集进行训练和评估。 更具体地说,我们使用 BLINK 数据集预训练检索器,使用 AIDA 数据集微调检索器并训练阅读器。
BLINK 数据集可以从 GENRE 仓库使用此
脚本 下载。
我们使用 blink-train-kilt.jsonl 和 blink-dev-kilt.jsonl 作为训练和验证数据集。
假设我们已在 data/blink 文件夹中下载了这两个文件,我们使用以下脚本将 BLINK 数据集转换为 ReLiK 格式:
# Train
python scripts/data/blink/preprocess_genre_blink.py \
data/blink/blink-train-kilt.jsonl \
data/blink/processed/blink-train-kilt-relik.jsonl
# Dev
python scripts/data/blink/preprocess_genre_blink.py \
data/blink/blink-dev-kilt.jsonl \
data/blink/processed/blink-dev-kilt-relik.jsonl
AIDA 数据集未公开提供,但我们提供了不含 text 字段的文件。您可以在 data/aida/processed 文件夹中找到 ReLiK 格式的文件。
我们使用的维基百科索引可从 此处 下载。
关系抽取
所有数据应具有以下结构:
{
"doc_id": int, # Unique identifier for the document
"doc_words: list[txt] # Tokenized text of the document
"doc_span_annotations": # Token level annotations of mentions (label is optional)
[
[start, end, label],
[start, end, label],
...
],
"doc_triplet_annotations": # Triplet annotations
[
{
"subject": [start, end, label], # label is optional
"relation": name, # type is optional
"object": [start, end, label], # label is optional
},
{
"subject": [start, end, label], # label is optional
"relation": name, # type is optional
"object": [start, end, label], # label is optional
},
]
}
对于关系抽取,我们提供了一个如何预处理来自 raw_nyt 的 NYT 数据集的示例,该数据集取自 CopyRE。将数据集下载到 data/raw_nyt,然后运行:
python scripts/data/nyt/preprocess_nyt.py data/raw_nyt data/nyt/processed/
请注意,为了进行公平比较,我们复现了先前工作中的预处理步骤,这会导致由于对实体跨度中重复的表面形式处理不当而产生重复的三元组。如果您希望正确地将原始数据解析为 ReLiK 格式,可以将标志设置为 --legacy-format False。只需注意,提供的 RE NYT 模型是在旧格式上训练的。
🦮 检索器
我们对检索器执行了两步训练过程。首先,我们使用 BLINK (Wu et al., 2019) 数据集对检索器进行“预训练”,然后使用 AIDA (Hoffart et al, 2011) 对其进行“微调”。
数据准备
检索器需要一个格式类似于 DPR: 一个 jsonl 文件,其中每一行是一个具有以下键的字典:
{
"question": "....",
"positive_ctxs": [{
"title": "...",
"text": "...."
}],
"negative_ctxs": [{
"title": "...",
"text": "...."
}],
"hard_negative_ctxs": [{
"title": "...",
"text": "...."
}]
}
检索器还需要一个索引来搜索文档。要索引的文档可以是 JSONL 文件或类似于 DPR:
jsonl: 每一行是一个 JSON 对象,包含以下键:id、text、metadatatsv: 每一行是一个制表符分隔的字符串,包含id和text列, 后跟任何其他列,这些列将存储在metadata字段中
jsonl 示例:
{
"id": "...",
"text": "...",
"metadata": ["{...}"]
},
...
tsv 示例:
id \t text \t any other column
...
实体链接
BLINK
一旦你获得了 ReLiK 格式的 BLINK 数据集,你可以使用以下脚本创建窗口:
# train
relik data create-windows \
data/blink/processed/blink-train-kilt-relik.jsonl \
data/blink/processed/blink-train-kilt-relik-windowed.jsonl
# dev
relik data create-windows \
data/blink/processed/blink-dev-kilt-relik.jsonl \
data/blink/processed/blink-dev-kilt-relik-windowed.jsonl
然后将其转换为 DPR 格式:
# train
relik data convert-to-dpr \
data/blink/processed/blink-train-kilt-relik-windowed.jsonl \
data/blink/processed/blink-train-kilt-relik-windowed-dpr.jsonl \
data/kb/wikipedia/documents.jsonl \
--title-map data/kb/wikipedia/title_map.json
# dev
relik data convert-to-dpr \
data/blink/processed/blink-dev-kilt-relik-windowed.jsonl \
data/blink/processed/blink-dev-kilt-relik-windowed-dpr.jsonl \
data/kb/wikipedia/documents.jsonl \
--title-map data/kb/wikipedia/title_map.json
AIDA
由于 AIDA 数据集未公开提供,我们可以提供 AIDA 数据集的 ReLiK 格式标注作为示例。
假设您拥有 data/aida 中的完整 AIDA 数据集,您可以将其转换为 ReLiK 格式,然后使用以下脚本创建窗口:
relik data create-windows \
data/aida/processed/aida-train-relik.jsonl \
data/aida/processed/aida-train-relik-windowed.jsonl
然后将其转换为 DPR 格式:
relik data convert-to-dpr \
data/aida/processed/aida-train-relik-windowed.jsonl \
data/aida/processed/aida-train-relik-windowed-dpr.jsonl \
data/kb/wikipedia/documents.jsonl \
--title-map data/kb/wikipedia/title_map.json
关系抽取
NYT
relik data create-windows \
data/data/processed/nyt/train.jsonl \
data/data/processed/nyt/train-windowed.jsonl \
--is-split-into-words \
--window-size none
然后将其转换为 DPR 格式:
relik data convert-to-dpr \
data/data/processed/nyt/train-windowed.jsonl \
data/data/processed/nyt/train-windowed-dpr.jsonl
训练模型
relik retriever train 命令可用于训练检索器。它需要以下参数:
config_path: 配置文件的路径。overrides: 对配置文件的覆盖列表,格式为key=value。
配置文件的示例可以在 relik/retriever/conf 文件夹中找到。
实体链接
relik/retriever/conf 中的配置文件是 pretrain_iterable_in_batch.yaml 和 finetune_iterable_in_batch.yaml,我们分别使用它们来预训练和微调检索器。
例如,要在 AIDA 数据集上训练检索器,您可以运行以下命令:
relik retriever train relik/retriever/conf/finetune_iterable_in_batch.yaml \
model.language_model=intfloat/e5-base-v2 \
data.train_dataset_path=data/aida/processed/aida-train-relik-windowed-dpr.jsonl \
data.val_dataset_path=data/aida/processed/aida-dev-relik-windowed-dpr.jsonl \
data.test_dataset_path=data/aida/processed/aida-test-relik-windowed-dpr.jsonl \
data.shared_params.documents_path=data/kb/wikipedia/documents.jsonl
关系抽取
relik/retriever/conf 中的配置文件为 finetune_nyt_iterable_in_batch.yaml,我们使用它来针对 NYT 数据集微调检索器。对于 cIE,我们复用前一步中从 BLINK 预训练得到的模型。
例如,要在 NYT 数据集上训练检索器,您可以运行以下命令:
relik retriever train relik/retriever/conf/finetune_nyt_iterable_in_batch.yaml \
model.language_model=intfloat/e5-base-v2 \
data.train_dataset_path=data/nyt/processed/nyt-train-relik-windowed-dpr.jsonl \
data.val_dataset_path=data/nyt/processed/nyt-dev-relik-windowed-dpr.jsonl \
data.test_dataset_path=data/nyt/processed/nyt-test-relik-windowed-dpr.jsonl
推理
通过向 relik retriever train 命令传递 train.only_test=True,您可以跳过训练,仅评估模型。
它还需要 PyTorch Lightning 检查点的路径以及用于评估的数据集。
relik retriever train relik/retriever/conf/finetune_iterable_in_batch.yaml \
train.only_test=True \
test_dataset_path=data/aida/processed/aida-test-relik-windowed-dpr.jsonl
model.checkpoint_path=path/to/checkpoint
检索器编码器可以通过以下命令从检查点中保存:
from relik.retriever.lightning_modules.pl_modules import GoldenRetrieverPLModule
checkpoint_path = "path/to/checkpoint"
retriever_folder = "path/to/retriever"
# If you want to push the model to the Hugging Face Hub set push_to_hub=True
push_to_hub = False
# If you want to push the model to the Hugging Face Hub set the repo_id
repo_id = "sapienzanlp/relik-retriever-e5-base-v2-aida-blink-encoder"
pl_module = GoldenRetrieverPLModule.load_from_checkpoint(checkpoint_path)
pl_module.model.save_pretrained(retriever_folder, push_to_hub=push_to_hub, repo_id=repo_id)
使用 push_to_hub=True 时,模型将被推送到 🤗 Hugging Face Hub,repo_id 作为模型将被推送到的仓库 id。
检索器需要一个索引来搜索文档。可以使用 relik retriever create-index 命令创建索引
relik retriever create-index --help
Usage: relik retriever build-index [OPTIONS] QUESTION_ENCODER_NAME_OR_PATH
DOCUMENT_PATH OUTPUT_FOLDER
╭─ Arguments ────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮
│ * question_encoder_name_or_path TEXT [default: None] [required] │
│ * document_path TEXT [default: None] [required] │
│ * output_folder TEXT [default: None] [required] │
╰────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
╭─ Options ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮
│ --document-file-type TEXT [default: jsonl] │
│ --passage-encoder-name-or-path TEXT [default: None] │
│ --indexer-class TEXT [default: relik.retriever.indexers.inmemory.InMemoryDocumentIndex] │
│ --batch-size INTEGER [default: 512] │
│ --num-workers INTEGER [default: 4] │
│ --passage-max-length INTEGER [default: 64] │
│ --device TEXT [default: cuda] │
│ --index-device TEXT [default: cpu] │
│ --precision TEXT [default: fp32] │
│ --push-to-hub --no-push-to-hub [default: no-push-to-hub] │
│ --repo-id TEXT [default: None] │
│ --help Show this message and exit. │
╰────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
借助编码器和索引,检索器可以从仓库 ID 或本地路径加载:
from relik.retriever import GoldenRetriever
encoder_name_or_path = "sapienzanlp/relik-retriever-e5-base-v2-aida-blink-encoder"
index_name_or_path = "sapienzanlp/relik-retriever-e5-base-v2-aida-blink-wikipedia-index"
retriever = GoldenRetriever(
question_encoder=encoder_name_or_path,
document_index=index_name_or_path,
device="cuda", # or "cpu"
precision="16", # or "32", "bf16"
index_device="cuda", # or "cpu"
index_precision="16", # or "32", "bf16"
)
然后它可以用于检索文档:
retriever.retrieve("Michael Jordan was one of the best players in the NBA.", top_k=100)
🤓 Reader
Reader 负责从一组候选项(例如,可能的实体或关系)中,从文档中提取实体和关系。
Reader 可以被训练用于跨度提取或三元组提取。
RelikReaderForSpanExtraction 用于跨度提取,即实体链接,而 RelikReaderForTripletExtraction 用于三元组提取,即关系抽取。
Data Preparation
Reader 需要我们在 Before You Start 章节中创建的窗口化数据集,并补充来自 Retriever 的候选项。
可以使用 relik retriever add-candidates 命令将候选项添加到数据集中。
relik retriever add-candidates --help
Usage: relik retriever add-candidates [OPTIONS] QUESTION_ENCODER_NAME_OR_PATH
DOCUMENT_NAME_OR_PATH INPUT_PATH
OUTPUT_PATH
╭─ Arguments ─────────────────────────────────────────────────────────────────────────────────────────────────╮
│ * question_encoder_name_or_path TEXT [default: None] [required] │
│ * document_name_or_path TEXT [default: None] [required] │
│ * input_path TEXT [default: None] [required] │
│ * output_path TEXT [default: None] [required] │
╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
╭─ Options ───────────────────────────────────────────────────────────────────────────────────────────────────╮
│ --passage-encoder-name-or-path TEXT [default: None] │
│ --relations BOOLEAN [default: False] │
│ --top-k INTEGER [default: 100] │
│ --batch-size INTEGER [default: 128] │
│ --num-workers INTEGER [default: 4] │
│ --device TEXT [default: cuda] │
│ --index-device TEXT [default: cpu] │
│ --precision TEXT [default: fp32] │
│ --use-doc-topics --no-use-doc-topics [default: no-use-doc-topics] │
│ --help Show this message and exit. │
╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
实体链接
我们需要使用之前训练好的检索器(Retriever),为每个窗口添加候选项,这些候选项将被阅读器(Reader)使用。以下是使用我们已在 Aida 训练集上训练好的检索器的一个示例:
relik retriever add-candidates sapienzanlp/relik-retriever-e5-base-v2-aida-blink-encoder sapienzanlp/relik-retriever-e5-base-v2-aida-blink-wikipedia-index data/aida/processed/aida-train-relik-windowed.jsonl data/aida/processed/aida-train-relik-windowed-candidates.jsonl
关系抽取
关系抽取的情况也是如此。如果你想使用我们训练好的检索器:
relik retriever add-candidates sapienzanlp/relik-retriever-small-nyt-question-encoder sapienzanlp/relik-retriever-small-nyt-document-index data/nyt/processed/nyt-train-relik-windowed.jsonl data/nyt/processed/nyt-train-relik-windowed-candidates.jsonl
训练模型
与检索器类似,可以使用 relik reader train 命令来训练检索器。它需要以下参数:
config_path: 配置文件的路径。overrides: 对配置文件的覆盖列表,格式为key=value。
配置文件的示例可以在 relik/reader/conf 文件夹中找到。
实体链接
relik/reader/conf 中的配置文件是 large.yaml 和 base.yaml,我们分别使用它们来训练大型和基础阅读器。
例如,要在 AIDA 数据集上训练大型阅读器,请运行:
relik reader train relik/reader/conf/large.yaml \
train_dataset_path=data/aida/processed/aida-train-relik-windowed-candidates.jsonl \
val_dataset_path=data/aida/processed/aida-dev-relik-windowed-candidates.jsonl \
test_dataset_path=data/aida/processed/aida-dev-relik-windowed-candidates.jsonl
关系抽取
relik/reader/conf 中的配置文件为 large_nyt.yaml、base_nyt.yaml 和 small_nyt.yaml,我们分别使用它们来训练大型、基础型和小型阅读器。
例如,要在 AIDA 数据集上训练大型阅读器,请运行:
relik reader train relik/reader/conf/large_nyt.yaml \
train_dataset_path=data/nyt/processed/nyt-train-relik-windowed-candidates.jsonl \
val_dataset_path=data/nyt/processed/nyt-dev-relik-windowed-candidates.jsonl \
test_dataset_path=data/nyt/processed/nyt-test-relik-windowed-candidates.jsonl
推理
读者可以使用以下命令从检查点中保存:
from relik.reader.lightning_modules.relik_reader_pl_module import RelikReaderPLModule
checkpoint_path = "path/to/checkpoint"
reader_folder = "path/to/reader"
# If you want to push the model to the Hugging Face Hub set push_to_hub=True
push_to_hub = False
# If you want to push the model to the Hugging Face Hub set the repo_id
repo_id = "sapienzanlp/relik-reader-deberta-v3-large-aida"
pl_model = RelikReaderPLModule.load_from_checkpoint(
trainer.checkpoint_callback.best_model_path
)
pl_model.relik_reader_core_model.save_pretrained(experiment_path, push_to_hub=push_to_hub, repo_id=repo_id)
使用 push_to_hub=True 时,模型将被推送到 🤗 Hugging Face Hub,repo_id 作为模型上传的仓库 ID。
阅读器可以从仓库 ID 或本地路径加载:
from relik.reader.pytorch_modules.span import RelikReaderForSpanExtraction
from relik.reader.pytorch_modules.triplet import RelikReaderForTripletExtraction
# the reader for span extraction
reader_span = RelikReaderForSpanExtraction(
"sapienzanlp/relik-reader-deberta-v3-large-aida"
)
# the reader for triplet extraction
reader_tripltes = RelikReaderForTripletExtraction(
"sapienzanlp/relik-reader-deberta-v3-large-nyt"
)
并用于提取实体和关系:
# an example of candidates for the reader
candidates = ["Michael Jordan", "NBA", "Chicago Bulls", "Basketball", "United States"]
reader_span.read("Michael Jordan was one of the best players in the NBA.", candidates=candidates)
📊 性能
实体链接
我们使用 GERBIL 评估了 ReLiK 在实体链接任务上的性能。下表展示了 ReLiK Large 和 Base 的结果(InKB Micro F1):
| 模型 | AIDA | MSNBC | Der | K50 | R128 | R500 | O15 | O16 | Tot | OOD | AIT (m:s) |
|---|---|---|---|---|---|---|---|---|---|---|---|
| GENRE | 83.7 | 73.7 | 54.1 | 60.7 | 46.7 | 40.3 | 56.1 | 50.0 | 58.2 | 54.5 | 38:00 |
| EntQA | 85.8 | 72.1 | 52.9 | 64.5 | 54.1 | 41.9 | 61.1 | 51.3 | 60.5 | 56.4 | 20:00 |
| ReLiKsmall | 82.2 | 72.7 | 55.6 | 68.3 | 48.0 | 42.3 | 62.7 | 53.6 | 60.7 | 57.6 | 00:29 |
| ReLiKBase | 85.3 | 72.3 | 55.6 | 68.0 | 48.1 | 41.6 | 62.5 | 52.3 | 60.7 | 57.2 | 00:29 |
| ReLiKLarge | 86.4 | 75.0 | 56.3 | 72.8 | 51.7 | 43.0 | 65.1 | 57.2 | 63.4 | 60.2 | 01:46 |
对比系统在域内 AIDA 测试集以及域外 MSNBC (MSN)、Derczynski (Der)、KORE50 (K50)、N3-Reuters-128 (R128)、 N3-RSS-500 (R500)、OKE-15 (O15) 和 OKE-16 (O16) 测试集上的评估结果(InKB Micro F1)。粗体表示最佳模型。 GENRE 使用了提及词典。 AIT 列显示了系统使用 NVIDIA RTX 4090 处理整个 AIDA 测试集所需的时间(格式为分:秒,m:s), 除了 EntQA 无法在 24GB 内存中运行,因此使用了 A100。
为了评估 ReLiK,我们使用以下步骤:
-
从此处下载 GERBIL 服务器。
-
启动 GERBIL 服务器:
cd gerbil && ./start.sh
- 启动以下服务:
cd gerbil-SpotWrapNifWS4Test && mvn clean -Dmaven.tomcat.port=1235 tomcat:run
- 启动 ReLiK 服务器以提供 GERBIL,并将模型名称作为参数(例如
sapienzanlp/relik-entity-linking-large):
python relik/reader/utils/gerbil.py --relik-model-name sapienzanlp/relik-entity-linking-large
- 打开 URL http://localhost:1234/gerbil 并:
- 选择 A2KB 作为实验类型
- 选择 "Ma - strong annotation match"
- 在 Name 字段中填写您希望赋予该实验的名称
- 在 URI 字段中填写:http://localhost:1235/gerbil-spotWrapNifWS4Test/myalgorithm
- 选择数据集(我们使用 AIDA-B, MSNBC, Der, K50, R128, R500, OKE15, OKE16)
- 最后,运行实验
Relation Extraction
下表显示了 ReLiK Large 在 NYT 数据集上的结果(Micro F1):
| Model | NYT | NYT (Pretr) | AIT (m:s) |
|---|---|---|---|
| REBEL | 93.1 | 93.4 | 01:45 |
| UiE | 93.5 | -- | -- |
| USM | 94.0 | 94.1 | -- |
| ReLiKLarge | 95.0 | 94.9 | 00:30 |
要评估 Relation Extraction,我们可以直接使用 reader 和脚本 relik/reader/trainer/predict_re.py,指向已检索候选项的文件。如果您想使用我们训练好的 Reader:
python relik/reader/trainer/predict_re.py --model_path sapienzanlp/relik-reader-deberta-v3-large-nyt --data_path /Users/perelluis/Documents/relik/data/debug/test.window.candidates.jsonl --is-eval
请注意,我们基于开发集计算预测关系的阈值。在评估时计算它,您可以运行以下命令:
python relik/reader/trainer/predict_re.py --model_path sapienzanlp/relik-reader-deberta-v3-large-nyt --data_path /Users/perelluis/Documents/relik/data/debug/dev.window.candidates.jsonl --is-eval --compute-threshold
💽 引用本工作
如果您使用了本工作的任何部分,请考虑按以下方式引用该论文:
@inproceedings{orlando-etal-2024-relik,
title = "Retrieve, Read and LinK: Fast and Accurate Entity Linking and Relation Extraction on an Academic Budget",
author = "Orlando, Riccardo and Huguet Cabot, Pere-Llu{\'\i}s and Barba, Edoardo and Navigli, Roberto",
booktitle = "Findings of the Association for Computational Linguistics: ACL 2024",
month = aug,
year = "2024",
address = "Bangkok, Thailand",
publisher = "Association for Computational Linguistics",
}
🪪 许可证
数据和软件采用 Creative Commons Attribution-NonCommercial-ShareAlike 4.0 许可。