dexml_lf-wikipedia-500k

Distilbert encoder models trained on Wikipedia tagging dataset (LF-Wikipedia-500K) using the DEXML (Dual Encoder for eXtreme Multi-Label classification, ICLR'24) method.

Inference Usage (Sentence-Transformers)

With sentence-transformers installed you can use this model as following:

from sentence_transformers import SentenceTransformer
sentences = ["This is an example sentence", "Each sentence is converted"]
model = SentenceTransformer('quicktensor/dexml_lf-wikipedia-500k')
embeddings = model.encode(sentences)
print(embeddings)

Usage (HuggingFace Transformers)

With huggingface transformers you only need to be a bit careful with how you pool the transformer output to get the embedding, you can use this model as following;

from transformers import AutoTokenizer, AutoModel

pooler = lambda x: x[:, 0, :] # Choose CLS token and normalize

sentences = ["This is an example sentence", "Each sentence is converted"]
tokenizer = AutoTokenizer.from_pretrained('quicktensor/dexml_lf-wikipedia-500k')
model = AutoModel.from_pretrained('quicktensor/dexml_lf-wikipedia-500k')

encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
with torch.no_grad():
    embeddings = pooler(model(**encoded_input))

print(embeddings)

Cite

If you found this model helpful, please cite our work as:

@InProceedings{DEXML,
  author    = "Gupta, N. and Khatri, D. and Rawat, A-S. and Bhojanapalli, S. and Jain, P. and Dhillon, I.",
  title     = "Dual-encoders for Extreme Multi-label Classification",
  booktitle = "International Conference on Learning Representations",
  month     = "May",
  year      = "2024"
}