ITADN
README.md
以下内容由 AI 翻译,如有问题请点此提交 issue 反馈

检索、阅读与链接:在学术预算下快速且准确的实体链接与关系抽取

Conference Paper arXiv

relik Hugging Face Collection Hugging Face Spaces Open In Colab

Lightning PyTorch Code style: black PyPi Version Release Version

一个极速且轻量的信息抽取模型,用于实体链接关系抽取

🛠️ 安装

从 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

使用可选依赖项安装,以便使用 FastAPIRay 提供模型服务。

pip install relik[serve]

从源码安装

git clone https://github.com/SapienzaNLP/relik.git
cd relik
pip install -e .[all]

🤖 模型

论文中的模型:

完整的模型列表可在 🤗 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.jsonlblink-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 对象,包含以下键:idtextmetadata
  • tsv: 每一行是一个制表符分隔的字符串,包含 idtext 列, 后跟任何其他列,这些列将存储在 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.yamlfinetune_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.yamlbase.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.yamlbase_nyt.yamlsmall_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):

模型AIDAMSNBCDerK50R128R500O15O16TotOODAIT (m:s)
GENRE83.773.754.160.746.740.356.150.058.254.538:00
EntQA85.872.152.964.554.141.961.151.360.556.420:00
ReLiKsmall82.272.755.668.348.042.362.753.660.757.600:29
ReLiKBase85.372.355.668.048.141.662.552.360.757.200:29
ReLiKLarge86.475.056.372.851.743.065.157.263.460.201: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,我们使用以下步骤:

  1. 此处下载 GERBIL 服务器。

  2. 启动 GERBIL 服务器:

cd gerbil && ./start.sh
  1. 启动以下服务:
cd gerbil-SpotWrapNifWS4Test && mvn clean -Dmaven.tomcat.port=1235 tomcat:run
  1. 启动 ReLiK 服务器以提供 GERBIL,并将模型名称作为参数(例如 sapienzanlp/relik-entity-linking-large):
python relik/reader/utils/gerbil.py --relik-model-name sapienzanlp/relik-entity-linking-large
  1. 打开 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):

ModelNYTNYT (Pretr)AIT (m:s)
REBEL93.193.401:45
UiE93.5----
USM94.094.1--
ReLiKLarge95.094.900: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 许可。