README.md

---
license: apache-2.0
language:
- en
library_name: transformers
tags:
- bert
- information retrieval
- learned sparse model
---

Paper: [DeeperImpact: Optimizing Sparse Learned Index Structures](https://arxiv.org/abs/2405.17093)

This repository contains the DeeperImpact model trained on the MS-MARCO passage dataset expanded using a [fine-tuned Llama 2 model](https://huggingface.co/soyuj/llama2-doc2query)
with hard negatives, distillation, and pre-trained CoCondenser model initialization.

The code to train and run inferences using DeeperImpact can be found in the [DeeperImpact Repo](https://github.com/basnetsoyuj/improving-learned-index).

Please refer to the following notebook to understand how to use the model: [inference_deeper_impact.ipynb](https://github.com/basnetsoyuj/improving-learned-index/blob/master/inference_deeper_impact.ipynb)

For running inference on a larger collection of documents, use the following command:

```bash
python -m src.deep_impact.index \
  --collection_path <expanded_collection.tsv> \
  --output_file_path <path> \
  --model_checkpoint_path soyuj/deeper-impact \
  --num_processes <n> \
  --process_batch_size <process_batch_size> \
  --model_batch_size <model_batch_size>
```

It distributes the inference across multiple GPUs in the machine. To manually set the GPUs, use `CUDA_VISIBLE_DEVICES` environment variable.

## ONNX

An ONNX export is available at [`onnx/model.onnx`](onnx/model.onnx) for inference with [ONNX Runtime](https://onnxruntime.ai/) — e.g. from Rust/C++/JS, or from Python without PyTorch.

|        | names                                          | dtype   | shape          |
|--------|------------------------------------------------|---------|----------------|
| inputs | `input_ids`, `attention_mask`, `token_type_ids` | int64   | `[batch, seq]` |
| output | `impact_scores`                                | float32 | `[batch, seq]` |

`impact_scores` is a per-subword-token score. A term's impact is the score at its **first** subword token — the same indexing as `DeepImpact.compute_term_impacts` (`##` continuation subwords are skipped; punctuation and terms past the 512-token window are dropped). Batch and sequence axes are dynamic.

The file was exported with [`src/deep_impact/scripts/export_onnx.py`](https://github.com/basnetsoyuj/DeeperImpact/blob/master/src/deep_impact/scripts/export_onnx.py) and matches the PyTorch model within `max |diff| ~ 6e-6`.