通过矢量相似性搜寻，能够在〜50ms 内响应〜640K 论文上的语义搜寻查问
Arxiv.org 大家肯定都不生疏，学习数据迷信的最佳办法之一是浏览 Arxiv.org 上的开源钻研论文。然而即便对于经验丰富的钻研人员来说，从大量的钻研论文中找出想读的内容也是十分不容易的。Connected 等论文之类的工具能够提供一些帮忙，然而它们依据论文之间共享的援用和参考书目来掂量相似性的，这当然十分的好，并且也很简略，然而文档中文本的语义含意也是一个掂量类似度十分重要的特色。
在这篇文章中，咱们将手动构建一个语义相似性搜索引擎，该引擎将单个论文作为“查问”输出，并查找 Top- K 的最相似论文。咱们测试，在单个笔记本电脑中对 Arxiv 语料库中的 640k 计算机科学论文进行查问的的提早 <50ms！如果你感兴趣，那么本文的次要内容总结如下：
- 设置环境并从 Kaggle 下载 ARXIV 数据
- 应用 dask 将数据加载到 Python 中
- 应用 MILVUS 矢量数据库进行语义相似性搜寻
本文中应用的技术不仅仅局限在科学论文，可用作建设任何 NLP 语义相似性搜索引擎的模板。惟一的区别是应用的预训练模型不同。
这篇文章应用 Kaggle 的 ARXIV 数据集是在 CC0：公共域许可证下公布的，所以请先浏览其应用受权的要求。
Milvus Vector 是一个矢量数据库，咱们也能够应用其余矢量数据库，如果应用其余库替换的话，有许多步骤完全相同并且改变并不简单。
设置环境并从 Kaggle 下载 ARXIV 数据。
Cornel University 已将整个 Arxiv 语料库上传到 Kaggle，并依据 CC0：公共畛域许可证取得许可。咱们能够应用 Kaggle API 间接下载数据集。
这里还要应用 Conda 环境，创立一个称为 Semantic_sibilarity 的环境。上面的步骤是创立必要的目录和 Conda 环境，装置所需的 Python 库，而后从 Kaggle 下载 ARXIV 数据集。
# Create the necessary directories
mkdir -p semantic_similarity/notebooks semantic_similarity/data semantic_similarity/milvus
# CD into the data directory
cd semantic_similarity/data
# Create and activate a conda environment
conda create -n semantic_similarity python=3.9
conda activate semantic_similarity
## Create Virtual Environment using venv if not using conda
# python -m venv semantic_similarity
# source semantic_similarity/bin/activate
# Pip install the necessary libraries
pip install jupyterlab kaggle matplotlib scikit-learn tqdm ipywidgets
pip install "dask[complete]" sentence-transformers
pip install pandas pyarrow pymilvus protobuf==3.20.0
# Download data using the kaggle API
kaggle datasets download -d Cornell-University/arxiv
# Unzip the data into the local directory
unzip arxiv.zip
# Delete the Zip file
rm arxiv.zip应用 dask 将数据加载到 Python 中
咱们从 Kaggle 下载的数据是一个 3.3GB JSON 文件，其中蕴含大概 200 万篇论文！为了无效地解决如此大的数据集，应用 PANDA 将整个数据集加载到内存中并不是一个好主见。为了解决这样大的数据，咱们抉择应用 DASK 将数据分为多个分区，并且仅将一些须要解决的分区加载到内存中。
Dask
Dask 是一个开源库，能够让咱们应用相似于 PANDA 的 API 进行并行计算。通过运行“pip install dask[complete]”在本地计算机上进行装置。装置实现后要导入必要的库。
import dask.bag as db
import json
from datetime import datetime
import time
data_path = '../data/arxiv-metadata-oai-snapshot.json'咱们将应用两个无效地解决大型 ARXIV JSON 文件的 DASK 的组件。
- Dask Bag：使咱们能够将 JSON 文件加载到固定大小的块中，并在每行数据上运行一些预处理性能
- DASK DATAFRAME：将 DASK Bag 转换为 DASK DATAFRAME，并能够用相似 Pandas 的 API 拜访
步骤 1：将 JSON 文件加载到 Dask Bag 中
将 JSON 文件加载到一个 Dask Bag 中，每个块的大小为 10MB。能够调整 blocksize 参数，管制每个块的大小。而后应用.map（）函数将 JSON.LOADS 函数利用于 Dask Bag 的每一行，将 JSON 字符串解析为 Python 字典。
# Read the file in blocks of 10MB and parse the JSON.
papers_db = db.read_text(data_path, blocksize="10MB").map(json.loads)
# Print the first row
papers_db.take(1)步骤 2：编写预处理辅助函数
从打印输出中能够看到每行蕴含与论文相干的几个元数据。让咱们编写三个辅助函数，能够帮忙咱们对数据集进行预处理。
v1_date（）：此函数是提取作者将论文的第一个版上传到 arxiv 的日期。咱们将将日期转换为 UNIX 工夫戳，并将其存储在该行中新的字段。
text_col（）：此函数是应用“[sep]”令牌组合“题目”和“摘要”字段，以便咱们能够将这些文本发送到 SPECTRE embedding 模型中。
filters（）：此函数过滤合乎某些条件的行，例如计算机科学类别中各个列和论文中的最大文本长度等等。
def v1_date(row):
"""
For each row in the dask bag,
find the date of the first version of the paper
and add it to the row as a new column
Args:
row: a row of the dask bag
Returns:
A row of the dask bag with added "unix_time" column
"""versions = row["versions"]
date = None
for version in versions:
if version["version"] == "v1":
date = datetime.strptime(version["created"], "%a, %d %b %Y %H:%M:%S %Z")
date = int(time.mktime(date.timetuple()))
row["unix_time"] = date
return row
def text_col(row):
"""It takes a row of a dataframe, adds a new column called'text'that is the concatenation of the'title'and'abstract' columns
Args:
row: the row of the dataframe
Returns:
A row with the text column added.
"""row["text"] = row["title"] +"[SEP]"+ row["abstract"]
return row
def filters(row):
"""
For each row in the dask bag, only keep the row if it meets the filter criteria
Args:
row: the row of the dataframe
Returns:
Boolean mask
"""return ((len(row["id"])<16) and
(len(row["categories"])<200) and
(len(row["title"])<4096) and
(len(row["abstract"])<65535) and
("cs." in row["categories"]) # Keep only CS papers
)步骤 3：在 Dask Bag 上运行预处理辅助函数
如下所示，咱们能够应用.map（）和.filter（）函数在 Dask Bag 的每一行上运行。因为 Dask 反对办法链，因而咱们能够仅保留一些必须的列，而后删除不须要的列。
# Specify columns to keep in the final table
cols_to_keep = ["id", "categories", "title", "abstract", "unix_time", "text"]
# Apply the pre-processing
papers_db = (papers_db.map(lambda row: v1_date(row))
.map(lambda row: text_col(row))
.map(
lambda row: {
key: value
for key, value in row.items()
if key in cols_to_keep
}
)
.filter(filters)
)
# Print the first row
papers_db.take(1)步骤 4：将 Dask Bag 转换为 DASK DATAFRAME
数据加载的最初一步是将 Dask Bag 转换为 DASK DATAFRAME，这样咱们能够应用相似 Pandas 的 API 进行拜访。
# Convert the Dask Bag to a Dask Dataframe
schema = {
"id": str,
"title": str,
"categories": str,
"abstract": str,
"unix_time": int,
"text": str,
}
papers_df = papers_db.to_dataframe(meta=schema)
# Display first 5 rows
papers_df.head()应用 MILVUS 矢量数据库进行语义相似性搜寻
Milvus 是最受欢迎的开源矢量数据库之一，所以咱们在本文中抉择应用它，并且咱们这里应用的是单机版，因为咱们只在本地机器上运行 Milvus。
步骤 1：本地装置 MILVUS 矢量数据库
应用 Docker 装置 Milvus Vector 数据库很简略，因而咱们首先须要装置 Docker。而后就是下载 Docker-compose.yml 并启动 Docker 容器，如下所示！MILVUS.IO 网站提供了许多其余抉择来装置 Milvus 单机版和 Milvus 群集版；如果须要在 Kubernetes 群集上装置或离线装置，请参考具体文档。
# CD into milvus directory
cd semantic_similarity/milvus
# Download the Standalone version of Milvus docker compose
wget https://github.com/milvus-io/milvus/releases/download/v2.1.0/milvus-standalone-docker-compose.yml -O ./docker-compose.yml
# Run the Milvus server docker container on your local
sudo docker-compose up -d步骤 2：创立一个 Milvus 汇合
咱们能够应用 Pymilvus 库与 Milvus Vector 数据库服务进行交互。emb_dim 参数是文本转换为嵌入的维度。在 SPECTRE 的状况下，嵌入维度为 768。
# Make sure a Milvus server is already running
from pymilvus import connections, utility
from pymilvus import Collection, CollectionSchema, FieldSchema, DataType
# Connect to Milvus server
connections.connect(alias="default", host="localhost", port="19530")
# Collection name
collection_name = "arxiv"
# Embedding size
emb_dim = 768
# # Check for existing collection and drop if exists
# if utility.has_collection(collection_name):
# print(utility.list_collections())
# utility.drop_collection(collection_name)Milvus 的汇合是相似于传统数据库中的表格。要创立一个汇合，首先须要指定汇合的模式。在本文示例中利用 Milvus 2.1 字符串索引和字段来存储与每篇论文相干的所有必要元数据。主键 idx 和其余字段 categories、title、abstract 是 VARCHAR 数据类型，而嵌入是蕴含 emb_dim 维度嵌入的 FLOAT_VECTOR 字段。Milvus 反对多种数据类型，如下所示。
# Create a schema for the collection
idx = FieldSchema(name="id", dtype=DataType.VARCHAR, is_primary=True, max_length=16)
categories = FieldSchema(name="categories", dtype=DataType.VARCHAR, max_length=200)
title = FieldSchema(name="title", dtype=DataType.VARCHAR, max_length=4096)
abstract = FieldSchema(name="abstract", dtype=DataType.VARCHAR, max_length=65535)
unix_time = FieldSchema(name="unix_time", dtype=DataType.INT64)
embedding = FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=emb_dim)
# Fields in the collection
fields = [idx, categories, title, abstract, unix_time, embedding]
schema = CollectionSchema(fields=fields, description="Semantic Similarity of Scientific Papers")
# Create a collection with the schema
collection = Collection(name=collection_name, schema=schema, using="default", shards_num=10)一旦创立了汇合，当初就能够将文本和向量插入到汇合中。
步骤 3: 遍历 Dask 分区，应用 SPECTER 进行文本嵌入，并将它们插入到 Milvus。
咱们须要将 Dask DATAFRAME 中的文本转换为嵌入向量来进行语义类似度搜寻。所以首先须要生成文本的嵌入。本文将应用名为 SPECTRE 的 SBERT 双编码器模型。
SPECTER : Scientific Paper Embeddings using Citation-informed TransformERs。
简略的说 SPECTER 是通过论文数据进行专门训练的模型，所以在选题分类、引文预测、科学论文举荐等方面的体现优于 SciBERT，这就是咱们抉择他的文章。
通过 Sentence Transformer 库，应用事后训练的 SPECTRE 模型非常简单。只须要一行代码就能够下载预训练的模型，咱们还编写了一个简略的辅助函数，将 Dask dataframe 分区的整个文本列转换为嵌入。
from sentence_transformers import SentenceTransformer
from tqdm import tqdm
# Scientific Papers SBERT Model
model = SentenceTransformer('allenai-specter')
def emb_gen(partition):
return model.encode(partition['text']).tolist()咱们能够应用 dask.map_partitions() API 将嵌入生成的函数利用到分区中的每一行，而后能够应用 collection.insert 将数据上传到 Milvus。
# Initialize
collection = Collection(collection_name)
for partition in tqdm(range(papers_df.npartitions)):
# Get the dask dataframe for the partition
subset_df = papers_df.get_partition(partition)
# Check if dataframe is empty
if len(subset_df.index) != 0:
# Metadata
data = [subset_df[col].values.compute().tolist()
for col in ["id", "categories", "title", "abstract", "unix_time"]
]
# Embeddings
data += [
subset_df
.map_partitions(emb_gen)
.compute()[0]
]
# Insert data
collection.insert(data)须要留神的是增加到数据变量中的列的程序必须与创立时定义的字段变量的程序雷同!
步骤 4: 对插入的数据将创立一个近似最近街坊 (ANN) 索引
在咱们将所有的嵌入插入到 Milvus 向量数据库后，还须要创立一个神经网络索引来放慢搜寻速度。在这个例子中，我应用的是 HNSW 索引，这是最快、最精确的 ANN 索引之一。无关 HNSW 指数及其参数的更多信息，请参阅 Milvus 文档。
# Add an ANN index to the collection
index_params = {
"metric_type": "L2",
"index_type": "HNSW",
"params": {"efConstruction": 128, "M": 8},
}
collection.create_index(field_name="embedding", index_params=index_params)步骤 5: 运行向量类似度搜寻查问!
实现了以上步骤当前就能够查问 Milvus 汇合中的数据了。首先加载汇合：
collection = Collection(collection_name)
collection.load()接下来，我创立了一个简略的辅助函数，它接管 query_text 并将其转换为 SPECTRE 嵌入，在 Milvus 汇合中执行 ANN 搜寻，并打印出后果。还有一些 search_params 能够管制搜寻的品质和速度，请参考 Milvus 文档。
def query_and_display(query_text, collection, num_results=10):
# Embed the Query Text
query_emb = [model.encode(query_text)]
# Search Params
search_params = {"metric_type": "L2", "params": {"ef": 128}}
# Search
query_start = datetime.now()
results = collection.search(
data=query_emb,
anns_field="embedding",
param=search_params,
limit=num_results,
expr=None,
output_fields=["title", "abstract"],
)
query_end = datetime.now()
# Print Results
print(f"Query Speed: {(query_end - query_start).total_seconds():.2f} s")
print("Results:")
for res in results[0]:
title = res.entity.get("title").replace("\n", "")
print(f"➡️ ID: {res.id}. L2 Distance: {res.distance:.2f}")
print(f"Title: {title}")
print(f"Abstract: {res.entity.get('abstract')}")咱们看看后果怎么样
# Query for papers that are similar to the SimCSE paper
title = "SimCSE: Simple Contrastive Learning of Sentence Embeddings"
abstract = """This paper presents SimCSE, a simple contrastive learning framework that greatly advances state-of-the-art sentence embeddings. We first describe an unsupervised approach, which takes an input sentence and predicts itself in a contrastive objective, with only standard dropout used as noise. This simple method works surprisingly well, performing on par with previous supervised counterparts. We find that dropout acts as minimal data augmentation, and removing it leads to a representation collapse. Then, we propose a supervised approach, which incorporates annotated pairs from natural language inference datasets into our contrastive learning framework by using"entailment"pairs as positives and"contradiction"pairs as hard negatives. We evaluate SimCSE on standard semantic textual similarity (STS) tasks, and our unsupervised and supervised models using BERT base achieve an average of 76.3% and 81.6% Spearman's correlation respectively, a 4.2% and 2.2% improvement compared to the previous best results. We also show -- both theoretically and empirically -- that the contrastive learning objective regularizes pre-trained embeddings'anisotropic space to be more uniform, and it better aligns positive pairs when supervised signals are available."""
query_text = f"{title}[SEP]{abstract}"
query_and_display(query_text, collection, num_results=10)如果不须要查问了，能够开释汇合来开释机器的内存。
collection.release()这在单机运行时是很好的办法，然而如果提供线上的服务则不要这样利用，因为每次加载都须要读取硬盘的数据，会很慢。
总结
在这篇文章中，咱们应用 SPECTRE 嵌入和 Milvus 向量数据库和几个简略的步骤中实现了一个可扩大的科学论文语义搜寻服务。这种办法在生产中可扩大到数亿甚至数十亿的数据。Milvus 在 30 毫秒内返回了前 10 个后果，这个速度对咱们来说还是十分不错的。
https://avoid.overfit.cn/post/36ad9ebf46ad43f78b84595e793e1a34
作者：Marie Stephen Leo