transformers safetensors bert fill-mask learned sparse retrieval passage-retrieval document-expansion bag-of-words en base_model:atomic-canyon/fermi-bert-1024 base_model:finetune:atomic-canyon/fermi-bert-1024 license:apache-2.0 autotrain_compatible endpoints_compatible region:us
---
base_model:
- atomic-canyon/fermi-bert-1024
language: en
license: apache-2.0
tags:
- learned sparse
- transformers
- retrieval
- passage-retrieval
- document-expansion
- bag-of-words
---

<br><br>
<p align="center">
<a href="https://atomic-canyon.com/"><svg id="Layer_1" data-name="Layer 1" xmlns="http://www.w3.org/2000/svg" width="450" viewBox="0 0 548.18 92.96"> <defs> <style> .cls-1 { stroke-width: 0px; } </style> </defs> <g> <path class="cls-1" d="m144.69,56.27h-16.94l-2.54,6.14h-7.06l14.26-32.63h7.71l14.31,32.63h-7.2l-2.54-6.14Zm-2.4-5.82l-6.05-14.59-6.05,14.59h12.09Z"/> <path class="cls-1" d="m171.91,62.4h-6.74v-26.31h-12.37v-6h31.48v6h-12.37v26.31Z"/> <path class="cls-1" d="m203.99,63.05c-10.15,0-17.12-6.92-17.12-16.76s7.02-16.85,17.26-16.85,17.08,6.92,17.08,16.76-7.06,16.85-17.22,16.85Zm.05-27.83c-6.09,0-10.29,4.57-10.29,10.99s4.25,11.08,10.34,11.08,10.29-4.57,10.29-10.99-4.29-11.08-10.34-11.08Z"/> <path class="cls-1" d="m261.22,37.57l-10.11,24.88h-6.46l-10.02-24.79-1.48,24.74h-6.51l2.12-32.31h9.32l9.88,24.6,9.97-24.6h9.28l2.03,32.31h-6.65l-1.38-24.83Z"/> <path class="cls-1" d="m276.67,62.4V30.09h6.74v32.31h-6.74Z"/> <path class="cls-1" d="m307.73,62.96c-10.48,0-17.86-6.46-17.86-16.71s7.75-16.71,17.72-16.71c4.43,0,8.12.88,11.59,2.31l-1.52,5.91c-2.95-1.29-6.23-2.22-9.79-2.22-6.55,0-11.12,4.34-11.12,10.62,0,6.6,4.62,10.8,11.45,10.8,3.18,0,6.37-.79,9.6-2.22l1.52,5.63c-3.6,1.71-7.57,2.58-11.59,2.58Z"/> <path class="cls-1" d="m351.39,62.96c-10.48,0-17.86-6.46-17.86-16.71s7.75-16.71,17.72-16.71c4.43,0,8.12.88,11.59,2.31l-1.52,5.91c-2.95-1.29-6.23-2.22-9.79-2.22-6.55,0-11.12,4.34-11.12,10.62,0,6.6,4.62,10.8,11.45,10.8,3.18,0,6.37-.79,9.6-2.22l1.52,5.63c-3.6,1.71-7.57,2.58-11.59,2.58Z"/> <path class="cls-1" d="m392.56,56.27h-16.94l-2.54,6.14h-7.06l14.26-32.63h7.71l14.31,32.63h-7.2l-2.54-6.14Zm-2.4-5.82l-6.05-14.59-6.05,14.59h12.09Z"/> <path class="cls-1" d="m413.92,39.88v22.52h-6.55V30.09h7.25l16.62,22.2v-22.2h6.51v32.31h-6.92l-16.89-22.52Z"/> <path class="cls-1" d="m462.61,62.4h-6.69v-12.14l-13.66-20.17h7.57l9.51,14.4,9.46-14.4h7.48l-13.66,20.12v12.19Z"/> <path class="cls-1" d="m494.32,63.05c-10.16,0-17.12-6.92-17.12-16.76s7.02-16.85,17.26-16.85,17.08,6.92,17.08,16.76-7.06,16.85-17.22,16.85Zm.05-27.83c-6.09,0-10.29,4.57-10.29,10.99s4.25,11.08,10.34,11.08,10.29-4.57,10.29-10.99-4.29-11.08-10.34-11.08Z"/> <path class="cls-1" d="m524.36,39.88v22.52h-6.55V30.09h7.25l16.62,22.2v-22.2h6.51v32.31h-6.92l-16.89-22.52Z"/> </g> <path class="cls-1" d="m66.85,4.93l-3.14,5.24s-15.72-8.38-34.06,0c0,0-23.06,9.43-22.53,38.25,0,0-.26,16.51,15.46,29.61,0,0,17.29,15.2,40.87,4.19l3.67,5.76s-22.53,13.62-48.73-4.19c0,0-17.55-13.36-18.34-35.63C.04,48.16-2.05,21.96,23.62,5.71c0,0,19.65-12.31,43.23-.79Z"/> <path class="cls-1" d="m70,13.84l3.67-5.76s17.29,11,18.86,30.39c0,0,7.07,26.46-18.6,46.37l-27.25-46.11-9.43,16.77h11l3.67,6.81h-26.2l20.96-36.68,28.3,48.21s11-6.81,12.05-27.77c0,0,1.05-20.44-17.03-32.23Z"/></svg></a>
</p>
<br><br>

# fermi-1024: Sparse Retrieval Model for Nuclear Power

This **sparse retrieval model** is optimized for nuclear-specific applications. It encodes both queries and documents into high-dimensional sparse vectors, where the non-zero dimensions correspond to specific tokens in the vocabulary, and their values indicate the relative importance of those tokens. 

The vocabulary, and thus the sparse embeddings, are based on a nuclear-specific tokenizer. For example, terms like "NRC" are represented as single tokens rather than being split into multiple tokens. This approach improves both accuracy and efficiency. To achieve this, we trained a nuclear-specific [BERT base model](https://huggingface.co/atomic-canyon/fermi-bert-1024).

### Specifications

- **Developed by:** [Atomic Canyon](https://atomic-canyon.com/)
- **Finetuned from model:** [fermi-bert-1024](https://huggingface.co/atomic-canyon/fermi-bert-1024)
- **Context Length:** 1024
- **Vocab Size:** 30522
- **License:** `Apache 2.0`

## Training

`fermi-1024` was trained on [MS MARCO Passage Dataset](https://microsoft.github.io/msmarco/) using the [LSR framework](https://github.com/thongnt99/learned-sparse-retrieval) using the teacher model [ms-marco-MiniLM-L-6-v2](https://huggingface.co/cross-encoder/ms-marco-MiniLM-L-6-v2). Trained on the Oak Ridge National Laboratory [Frontier supercomputer](https://www.olcf.ornl.gov/frontier/) using MI250X AMD GPUs.

## Evaluation

The sparse embedding model was primarily evaluated for its effectiveness in information retrieval within the nuclear energy domain. Due to the absence of domain-specific benchmarks, we developed [FermiBench](https://huggingface.co/datasets/atomic-canyon/FermiBench) to assess the model’s performance on nuclear-related texts. In addition, the model was tested on the MS MARCO dev split and the BEIR benchmark to ensure broader applicability. The model demonstrates strong retrieval capabilities, particularly in handling nuclear-specific jargon and documents.

Although there are standard benchmarks and tooling for evaluating dense embedding models, we found no open, standardized tooling for evaluating sparse embedding models. To support the community, we are [releasing our benchmark tooling](https://github.com/atomic-canyon/fermi), built on top of [BEIR](https://github.com/beir-cellar/beir) and [pyserini](https://github.com/castorini/pyserini). All evaluation numbers were produced with that tool and should therefore be reproducible.

| Model                             | FermiBench NDCG@10 | FermiBench FLOPS | MSMarco Dev NDCG@10 | BEIR* NDCG@10 | BEIR*  FLOPS | 
| --------------------------------- | ------------------ | ---------------- | ------------------- | ------------- | ------------ |
| fermi-512                         | 0.74               | 7.07             | 0.45                | 0.46          | 9.14         |
| fermi-1024                        | 0.72               | 4.75             | 0.44                | 0.46          | 7.5          |
| splade-cocondenser-ensembledistil | 0.64               | 12.9             | 0.45                | 0.46          | 12.4         |

\* BEIR benchmark was a subset containng trec-covid, nfcorpus, arguana, scidocs, scifact.

### Efficiency

Given the massive scale of documentation in nuclear energy, efficiency is crucial. Our model addresses this in several ways:

- Our 1024-length embedding model reduces the number of required embeddings by half, significantly lowering computational costs.
- The custom tokenizer, designed for nuclear-specific jargon, encodes documents and queries using fewer tokens, improving computational efficiency.
- Additionally, our models produce sparser vectors, reducing FLOPs and, as a secondary benefit, lowering storage requirements for indexing.

## Usage

```python
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))
```

# Acknowledgement

This research used resources of the Oak Ridge Leadership Computing Facility, which is a DOE Office of Science User Facility supported under Contract DE-AC05-00OR22725.