STonKGs

STonKGs 是一个可以共同训练生物医学文本和知识图谱的高级 Transformer。这种多模态 Transformer 将来自 KG 的结构化信息与无结构化文本数据相结合，以学习联合表示。虽然我们在一个生物医学知识图谱（即来自 INDRA）上展示了 STonKGs，但该模型可以应用于其他领域。在以下各节中，我们将描述运行在任意给定数据集上训练模型的必要脚本。

💪 入门指南

数据格式

由于 STonKGs 在文本和 KG 数据上运行，因此预期相应数据文件包含两个模态的列。更具体地说，预期数据格式是 pandas dataframe（或用于预训练脚本的 pickled pandas dataframe），其中每行包含一个文本-三元组对。以下列是预期的

• source：给定文本-三元组对中的三元组源节点
• target：给定文本-三元组对中的三元组目标节点
• evidence：给定文本-三元组对中的文本
• （可选）class：在微调任务中给定文本-三元组对的类别标签（不适用于预训练过程）

请注意，源节点和目标节点都必须以生物表达语言（BEL）格式表示，更具体地说，它们需要包含在 INDRA KG 中。有关 BEL 格式的更多详细信息，请参阅例如 INDRA 的 BEL 处理器文档 和 PyBEL。

预训练 STonKGs

一旦您将 STonKGs 作为 Python 包安装（见下文），您可以通过运行以下命令开始对您的数据集进行 STonKGs 训练

$ python3 -m stonkgs.models.stonkgs_pretraining

可以通过修改 pretrain_stonkgs 方法中的参数轻松修改模型的配置。唯一需要更改的参数是 PRETRAINING_PREPROCESSED_POSITIVE_DF_PATH，它应指向您的数据集。

在 INDRA KG 上下载预训练的 STonKGs 模型

为了可能进行未来的适应，如在其他 KG 上进行进一步预训练，我们发布了在 INDRA KG 上预训练的 STonKGs 模型。STonKGs_150k 和 STonKGs_300k 都可以通过 Hugging Face 的模型库访问。

下载并初始化预训练的 STonKGs 模型最简单的方法是使用 from_default_pretrained() 类方法（默认为 STonKGs_150k）

from stonkgs import STonKGsForPreTraining

# Download the model from the model hub and initialize it for pre-training 
# using from_default_pretrained
stonkgs_pretraining = STonKGsForPreTraining.from_default_pretrained()

或者，由于我们的代码基于 Hugging Face 的 transformers 包，可以使用 .from_pretrained() 函数轻松下载和初始化预训练模型

from stonkgs import STonKGsForPreTraining

# Download the model from the model hub and initialize it for pre-training 
# using from_pretrained
stonkgs_pretraining = STonKGsForPreTraining.from_pretrained(
    'stonkgs/stonkgs-150k',
)

提取嵌入

预训练的 STonKGs 模型（或您自己的 STonKGs 变体）的学习嵌入可以通过两个简单的步骤提取。首先，需要使用 preprocess_file_for_embeddings 函数对包含文本-三元组的给定数据集（pandas DataFrame，见 数据格式）进行预处理。然后，可以使用预处理数据和 get_stonkgs_embeddings 函数获得学习嵌入

import pandas as pd

from stonkgs import get_stonkgs_embeddings, preprocess_df_for_embeddings

# Generate some example data
# Note that the evidence sentences are typically longer than in this example data
rows = [
    [
        "p(HGNC:1748 ! CDH1)",
        "p(HGNC:2515 ! CTNND1)",
        "Some example sentence about CDH1 and CTNND1.",
    ],
    [
        "p(HGNC:6871 ! MAPK1)",
        "p(HGNC:6018 ! IL6)",
        "Another example about some interaction between MAPK and IL6.",
    ],
    [
        "p(HGNC:3229 ! EGF)",
        "p(HGNC:4066 ! GAB1)",
        "One last example in which Gab1 and EGF are mentioned.",
    ],
]
example_df = pd.DataFrame(rows, columns=["source", "target", "evidence"])

# 1. Preprocess the text-triple data for embedding extraction
preprocessed_df_for_embeddings = preprocess_df_for_embeddings(example_df)

# 2. Extract the embeddings 
embedding_df = get_stonkgs_embeddings(preprocessed_df_for_embeddings)

微调 STonKGs

在原始六个分类任务上微调 STonKGs 最直接的方法是运行微调脚本（请注意，此脚本假定您已指定 mlflow 日志记录器，例如使用 --logging_dir 参数）

$ python3 -m stonkgs.models.stonkgs_finetuning

此外，在您自己的代码中使用 STonKGs 进行您自己的微调任务（即序列分类任务）与初始化预训练模型一样简单

from stonkgs import STonKGsForSequenceClassification

# Download the model from the model hub and initialize it for fine-tuning
stonkgs_model_finetuning = STonKGsForSequenceClassification.from_default_pretrained(
    num_labels=number_of_labels_in_your_task,
)

# Initialize a Trainer based on the training dataset
trainer = Trainer(
    model=model,
    args=some_previously_defined_training_args,
    train_dataset=some_previously_defined_finetuning_data,
)

# Fine-tune the model to the moon 
trainer.train()

使用 STonKGs 进行推理

您可以根据以下两种方式生成未见过的文本三元组对的新预测（只要节点包含在INDRA知识图谱中）：1) 用于基准的微调模型；2) 您自己的微调模型。为了做到这一点，您首先需要加载/初始化微调模型。

from stonkgs.api import get_species_model, infer

model = get_species_model()

# Next, you want to use that model on your dataframe (consisting of at least source, target
# and evidence columns, see **Data Format**) to generate the class probabilities for each
# text-triple pair belonging to each of the specified classes in the respective fine-tuning task:
example_data = ...

# See Extracting Embeddings for the initialization of the example data
# This returns both the raw (transformers) PredictionOutput as well as the class probabilities 
# for each text-triple pair
raw_results, probabilities = infer(model, example_data)

⬇️ 安装

可以从PyPI安装最新版本。

$ pip install stonkgs

可以直接从GitHub安装最新代码和数据。

$ pip install git+https://github.com/stonkgs/stonkgs.git

以开发模式安装，请使用以下命令

$ git clone git+https://github.com/stonkgs/stonkgs.git
$ cd stonkgs
$ pip install -e .

警告：由于stellargraph 目前不支持Python 3.9，此软件只能安装在Python 3.8上。

工件

预训练模型托管在HuggingFace。微调模型托管在STonKGs社区页面的Zenodo上，以及其他工件（node2vec嵌入、随机游走等）。

致谢

⚖️ 许可证

本软件包中的代码根据MIT许可证授权。

📖 引用

Balabin H.，Hoyt C.T.，Birkenbihl C.，Gyori B.M.，Bachman J.A.，Komdaullil A.T.，Plöger P.G.，Hofmann-Apitius M.，Domingo-Fernández D. STonKGs: A Sophisticated Transformer Trained on Biomedical Text and Knowledge Graphs (2021)，bioRxiv，TODO。

🎁 支持

该项目得到了几个组织的支持（按字母顺序排列）

• Fraunhofer机器学习中心
• 哈佛治疗科学计划 - 系统药理学实验室

💰 资金

该项目得到了以下资助

🍪 Cookiecutter

本软件包使用@audreyfeldroy的cookiecutter软件包以及@cthoyt的cookiecutter-snekpack模板创建。

🛠️ 开发

README的最后一部分是如果您想通过代码贡献来参与其中。

❓ 测试

在克隆存储库并使用pip install tox安装tox后，可以重复运行tests/文件夹中的单元测试。

$ tox

此外，这些测试会随着每个提交在GitHub动作中自动重新运行。

📦 发布版本

在开发模式下安装软件包并使用pip install tox安装tox后，创建新版本的命令包含在tox.ini中的finish环境中。在shell中运行以下命令：

$ tox -e finish

此脚本执行以下操作：

1. 使用BumpVersion将setup.cfg和src/stonkgs/version.py中的版本号切换为不带-dev后缀。
2. 将代码打包成tar存档和wheel格式。
3. 使用twine上传到PyPI。确保已经配置了.pypirc文件以避免在此步骤中需要手动输入。
4. 推送到GitHub。您需要创建一个发布版本，与版本号升级的提交一起进行。
5. 将版本升级到下一个补丁。如果您进行了重大更改并希望通过次要版本升级版本，可以在之后使用tox -e bumpversion minor。