zuliaio/specter2-proximity-onnx

SPECTER2 proximity (ONNX)

A single-file ONNX export of SPECTER2 with the proximity adapter activated. SPECTER2 produces scientific document embeddings trained on citation links (Singh et al., 2022). The proximity adapter is the variant intended for document-to-document similarity and nearest-neighbor retrieval.

This repository exists so the model can be loaded by ONNX Runtime without the adapters Python library. It is used by Zulia as KnownEmbeddingModel.SPECTER2.

Reproducing

export_specter2_onnx.py in this repository is the exact script that produced model.onnx. It loads the base model and adapter from Hugging Face, traces the graph, checks the ONNX output against PyTorch and writes this folder:

pip install torch --index-url https://download.pytorch.org/whl/cpu
pip install adapters onnx onnxruntime huggingface_hub
python export_specter2_onnx.py --out ./specter2-proximity-onnx

Modifications from the upstream model

• The Pfeiffer bottleneck adapter from allenai/specter2 was activated on allenai/specter2_base and the complete forward pass was traced to ONNX (opset 17, fp32). The adapter layers are part of the graph.
• The graph takes input_ids, attention_mask and token_type_ids and returns last_hidden_state. Pooling is not part of the graph. Apply CLS pooling (take position 0) and L2 normalize.
• No weights were changed. Outputs match the PyTorch model to within floating point tolerance.

Usage

SPECTER2 expects the paper title and abstract joined by the tokenizer separator token: title[SEP]abstract. Maximum sequence length is 512 tokens. Embedding dimension is 768.

import numpy as np
import onnxruntime as ort
from transformers import AutoTokenizer

tok = AutoTokenizer.from_pretrained("zuliaio/specter2-proximity-onnx")
sess = ort.InferenceSession("model.onnx")
enc = tok(["Title[SEP]Abstract text"], padding=True, truncation=True, max_length=512, return_tensors="np")
hidden = sess.run(None, dict(enc))[0]
emb = hidden[:, 0, :]
emb = emb / np.linalg.norm(emb, axis=1, keepdims=True)

License and attribution

Apache License 2.0, unchanged from the upstream base model and adapter released by the Allen Institute for AI. See the SPECTER2 repository and SciRepEval. If you use this model, please cite:

@inproceedings{Singh2022SciRepEvalAM,
  title={SciRepEval: A Multi-Format Benchmark for Scientific Document Representations},
  author={Amanpreet Singh and Mike D'Arcy and Arman Cohan and Doug Downey and Sergey Feldman},
  booktitle={Conference on Empirical Methods in Natural Language Processing},
  year={2022},
  url={https://api.semanticscholar.org/CorpusID:254018137}
}
@inproceedings{specter2020cohan,
  title={SPECTER: Document-level Representation Learning using Citation-informed Transformers},
  author={Arman Cohan and Sergey Feldman and Iz Beltagy and Doug Downey and Daniel S. Weld},
  booktitle={ACL},
  year={2020}
}