🚀 fermi-1024：用于核能的稀疏检索模型

这是一个针对核能特定应用优化的稀疏检索模型。它将查询和文档都编码为高维稀疏向量，其中非零维度对应词汇表中的特定标记，其值表示这些标记的相对重要性。

词汇表以及由此产生的稀疏嵌入基于核能特定的分词器。例如，像 “NRC” 这样的术语被表示为单个标记，而不是拆分为多个标记。这种方法提高了准确性和效率。为了实现这一点，我们训练了一个核能特定的 BERT 基础模型。

🚀 快速开始

本模型是专为核能特定应用优化的稀疏检索模型，它能将查询和文档编码为高维稀疏向量，提升检索的准确性与效率。

✨ 主要特性

• 特定领域优化：针对核能领域应用进行优化，能更好地处理核能相关的查询和文档。
• 高效编码：使用核能特定的分词器，将查询和文档编码为高维稀疏向量，提高了准确性和效率。
• 节省计算资源：1024 长度的嵌入模型将所需嵌入数量减少一半，降低了计算成本；自定义分词器使用更少的标记进行编码，提高了计算效率；模型生成的向量更稀疏，减少了浮点运算次数并降低了索引存储需求。

📦 安装指南

文档未提及具体安装步骤，故跳过该部分内容。

💻 使用示例

基础用法

import itertools
import torch
from transformers import AutoModelForMaskedLM, AutoTokenizer


# get sparse vector from dense vectors with shape batch_size * seq_len * vocab_size
def get_sparse_vector(feature, output):
    values, _ = torch.max(output*feature["attention_mask"].unsqueeze(-1), dim=1)
    values = torch.log(1 + torch.relu(values))
    values[:,special_token_ids] = 0
    return values
    
# transform the sparse vector to a dict of (token, weight)
def transform_sparse_vector_to_dict(sparse_vector):
    sample_indices,token_indices=torch.nonzero(sparse_vector,as_tuple=True)
    non_zero_values = sparse_vector[(sample_indices,token_indices)].tolist()
    number_of_tokens_for_each_sample = torch.bincount(sample_indices).cpu().tolist()
    tokens = [id_to_token[_id] for _id in token_indices.tolist()]

    output = []
    end_idxs = list(itertools.accumulate([0]+number_of_tokens_for_each_sample))
    for i in range(len(end_idxs)-1):
        token_strings = tokens[end_idxs[i]:end_idxs[i+1]]
        weights = non_zero_values[end_idxs[i]:end_idxs[i+1]]
        output.append(dict(zip(token_strings, weights)))
    return output
    

# load the model
model = AutoModelForMaskedLM.from_pretrained("atomic-canyon/fermi-1024")
tokenizer = AutoTokenizer.from_pretrained("atomic-canyon/fermi-1024")

# set the special tokens and id_to_token transform for post-process
special_token_ids = [tokenizer.vocab[token] for token in tokenizer.special_tokens_map.values()]
id_to_token = [""] * tokenizer.vocab_size
for token, _id in tokenizer.vocab.items():
    id_to_token[_id] = token

query = "What is the maximum heat load per spent fuel assembly for the EOS-37PTH?"
document = "For the EOS-37PTH DSC, add two new heat load zone configurations (HLZCs) for the EOS37PTH for higher heat load assemblies, up to 3.5 kW/assembly, that also allow for damaged and failed fuel storage."

# encode the query & document
feature = tokenizer([query, document], padding=True, truncation=True, return_tensors='pt', return_token_type_ids=False)
output = model(**feature)[0]
sparse_vector = get_sparse_vector(feature, output)

# get similarity score
sim_score = torch.matmul(sparse_vector[0],sparse_vector[1])
print(sim_score)


query_token_weight, document_query_token_weight = transform_sparse_vector_to_dict(sparse_vector)
for token in sorted(query_token_weight, key=lambda x:query_token_weight[x], reverse=True):
    if token in document_query_token_weight:
        print("score in query: %.4f, score in document: %.4f, token: %s"%(query_token_weight[token],document_query_token_weight[token],token))

📚 详细文档

规格参数

训练情况

fermi-1024 模型在 MS MARCO Passage 数据集 上进行训练，使用 LSR 框架，教师模型为 ms-marco-MiniLM-L-6-v2。训练在橡树岭国家实验室的 Frontier 超级计算机 上使用 MI250X AMD GPU 进行。

评估情况

该稀疏嵌入模型主要针对核能领域的信息检索效果进行评估。由于缺乏特定领域的基准测试，我们开发了 FermiBench 来评估模型在核能相关文本上的性能。此外，模型还在 MS MARCO 开发集和 BEIR 基准测试上进行了测试，以确保更广泛的适用性。模型展现出强大的检索能力，尤其在处理核能特定的行话和文档方面表现出色。

虽然有评估密集嵌入模型的标准基准和工具，但我们未找到用于评估稀疏嵌入模型的开放、标准化工具。为了支持社区，我们正在 发布我们的基准测试工具，该工具基于 BEIR 和 pyserini 构建。所有评估数据均使用该工具生成，因此应该是可复现的。

* BEIR 基准测试是一个包含 trec-covid、nfcorpus、arguana、scidocs、scifact 的子集。

效率优势

鉴于核能领域文档的大规模特性，效率至关重要。我们的模型通过以下几种方式解决了这一问题：

• 1024 长度的嵌入模型将所需嵌入数量减少一半，显著降低了计算成本。
• 自定义分词器专为核能特定行话设计，使用更少的标记对文档和查询进行编码，提高了计算效率。
• 模型生成的向量更稀疏，减少了浮点运算次数，同时还降低了索引存储需求。

🔧 技术细节

本模型将查询和文档编码为高维稀疏向量，其中非零维度对应词汇表中的特定标记，其值表示这些标记的相对重要性。词汇表和稀疏嵌入基于核能特定的分词器，能够更好地处理核能领域的专业术语。模型在训练过程中使用了核能特定的 BERT 基础模型，并在相关数据集上进行训练和评估，以确保在核能领域的准确性和效率。

📄 许可证

本项目采用 Apache 2.0 许可证。

致谢

本研究使用了橡树岭领导力计算设施的资源，该设施是美国能源部科学办公室的用户设施，由合同 DE - AC05 - 00OR22725 支持。