INF Word-level Sparse Embedding (INF-WSE)
INF-WSE is a series of word-level sparse embedding models developed by INF TECH. These models are optimized to generate sparse, high-dimensional text embeddings that excel in capturing the most relevant information for search and retrieval, particularly in Chinese text.
Key Features:
- Optimized for Retrieval: INF-WSE is designed with retrieval tasks in mind. The sparse embeddings enable efficient matching between queries and documents, making it highly effective for semantic search, ranking, and information retrieval scenarios where speed and accuracy are critical.
- Word-level Sparse Embeddings: The model generates sparse representations at the word level, capturing essential semantic details that help improve the relevance of search results. This is particularly useful for Chinese language retrieval tasks, where word segmentation can significantly impact performance.
- Sparse Representation for Efficiency: Unlike dense embeddings that have a fixed number of dimensions, INF-WSE produces sparse embeddings where the dimensionality matches the vocabulary size. Most dimensions are set to zero, focusing only on the most significant terms. This sparsity reduces the computational load, enabling faster retrieval without compromising on precision.
Usage
Transformers
Infer embeddings
import torch
from transformers import AutoTokenizer, AutoModel
queries = ['电脑一体机由什么构成?', '什么是掌上电脑?']
documents = [
'电脑一体机,是由一台显示器、一个电脑键盘和一个鼠标组成的电脑。',
'掌上电脑是一种运行在嵌入式操作系统和内嵌式应用软件之上的、小巧、轻便、易带、实用、价廉的手持式计算设备。',
]
input_texts = queries + documents
tokenizer = AutoTokenizer.from_pretrained("infly/inf-wse-v1-base-zh", trust_remote_code=True, use_fast=False) # Fast tokenizer has not been supported yet
model = AutoModel.from_pretrained("infly/inf-wse-v1-base-zh", trust_remote_code=True)
model.eval()
max_length = 512
input_batch = tokenizer(input_texts, padding=True, max_length=max_length, truncation=True, return_tensors="pt")
with torch.no_grad():
embeddings = model(input_batch['input_ids'], input_batch['attention_mask'], return_sparse=False) # if return_sparse=True, return sparse tensor, else return dense tensor
scores = embeddings[:2] @ embeddings[2:].T
print(scores.tolist())
# [[21.224790573120117, 4.520412921905518], [10.290857315063477, 19.359437942504883]]
Convert embeddings to lexical weights
from collections import OrderedDict
def convert_embeddings_to_weights(embeddings, tokenizer):
values, indices = torch.sort(embeddings, dim=-1, descending=True)
token2weight = []
for i in range(embeddings.size(0)):
token2weight.append(OrderedDict())
non_zero_mask = values[i] != 0
tokens = tokenizer.convert_ids_to_tokens(indices[i][non_zero_mask])
weights = values[i][non_zero_mask].tolist()
for token, weight in zip(tokens, weights):
token2weight[i][token] = weight
return token2weight
token2weight = convert_embeddings_to_weights(embeddings, tokenizer)
print(token2weight[1])
# OrderedDict([('掌上', 3.4572525024414062), ('电脑', 2.6253132820129395), ('是', 2.0787220001220703), ('什么', 1.2899624109268188)])
Evaluation
C-MTEB Retrieval task
(Chinese Massive Text Embedding Benchmark)
Metric: nDCG@10
Model Name | Max Length | Average | Cmedqa | Covid | Du | Ecom | Medical | MMarco | T2 | Video |
---|---|---|---|---|---|---|---|---|---|---|
BM25-zh | - | 50.37 | 13.70 | 86.58 | 57.13 | 44.04 | 32.08 | 48.31 | 60.48 | 60.64 |
bge-m3-sparse | 512 | 57.00 | 24.50 | 76.09 | 71.51 | 50.49 | 43.93 | 59.28 | 71.76 | 58.43 |
inf-wse-v1-base-zh | 512 | 61.16 | 20.51 | 76.41 | 79.84 | 56.78 | 46.24 | 66.40 | 76.50 | 68.57 |
All results, except for BM25, are measured by building the sparse index via Qdrant.
- Downloads last month
- 12
This model does not have enough activity to be deployed to Inference API (serverless) yet. Increase its social
visibility and check back later, or deploy to Inference Endpoints (dedicated)
instead.
Model tree for infly/inf-wse-v1-base-zh
Base model
junnyu/roformer_chinese_base