license: apache-2.0
language:
- en
library_name: transformers
tags:
- bert
- information retrieval
- learned sparse model
Paper: DeeperImpact: Optimizing Sparse Learned Index Structures
This repository contains the DeeperImpact model trained on the MS-MARCO passage dataset expanded using a fine-tuned Llama 2 model 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.
Please refer to the following notebook to understand how to use the model: inference_deeper_impact.ipynb
For running inference on a larger collection of documents, use the following command:
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 for inference with ONNX Runtime — 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 and matches the PyTorch model within max |diff| ~ 6e-6.