Views
No views yet
1# Requires transformers>=4.36.0
2
3import torch.nn.functional as F
4from transformers import AutoModel, AutoTokenizer
5
6input_texts = [
7 "what is the capital of China?",
8 "how to implement quick sort in python?",
9 "北京",
10 "快排算法介绍"
11]
12
13model_name_or_path = 'Alibaba-NLP/gte-multilingual-base'
14tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
15model = AutoModel.from_pretrained(model_name_or_path, trust_remote_code=True)
16
17# Tokenize the input texts
18batch_dict = tokenizer(input_texts, max_length=8192, padding=True, truncation=True, return_tensors='pt')
19
20outputs = model(**batch_dict)
21
22dimension=768 # The output dimension of the output embedding, should be in [128, 768]
23embeddings = outputs.last_hidden_state[:, 0][:dimension]
24
25embeddings = F.normalize(embeddings, p=2, dim=1)
26scores = (embeddings[:1] @ embeddings[1:].T) * 100
27print(scores.tolist())
28
29# [[0.3016996383666992, 0.7503870129585266, 0.3203084468841553]]1# Requires sentence-transformers>=3.0.0
2
3from sentence_transformers import SentenceTransformer
4
5input_texts = [
6 "what is the capital of China?",
7 "how to implement quick sort in python?",
8 "北京",
9 "快排算法介绍"
10]
11
12model_name_or_path="Alibaba-NLP/gte-multilingual-base"
13model = SentenceTransformer(model_name_or_path, trust_remote_code=True)
14embeddings = model.encode(input_texts, normalize_embeddings=True) # embeddings.shape (4, 768)
15
16# sim scores
17scores = model.similarity(embeddings[:1], embeddings[1:])
18
19print(scores.tolist())
20# [[0.301699697971344, 0.7503870129585266, 0.32030850648880005]]docker run --gpus all -v $PWD/data:/app/.cache -p "7997":"7997" \
michaelf34/infinity:0.0.69 \
v2 --model-id Alibaba-NLP/gte-multilingual-base --revision "main" --dtype float16 --batch-size 32 --device cuda --engine torch --port 79971docker run --platform linux/amd64 \
2 -p 8080:80 \
3 -v $PWD/data:/data \
4 --pull always \
5 ghcr.io/huggingface/text-embeddings-inference:cpu-1.7 \
6 --model-id Alibaba-NLP/gte-multilingual-base \
7 --dtype float16docker run --gpus all \
-p 8080:80 \
-v $PWD/data:/data \
--pull always \
ghcr.io/huggingface/text-embeddings-inference:1.7 \
--model-id Alibaba-NLP/gte-multilingual-base \
--dtype float16v1/embeddings route (more information about the OpenAI Embeddings API):1curl https://0.0.0.0:8080/v1/embeddings \
2 -H "Content-Type: application/json" \
3 -d '{
4 "input": [
5 "what is the capital of China?",
6 "how to implement quick sort in python?",
7 "北京",
8 "快排算法介绍"
9 ],
10 "model": "Alibaba-NLP/gte-multilingual-base",
11 "encoding_format": "float"
12 }'1# You can find the script gte_embedding.py in https://huggingface.co/Alibaba-NLP/gte-multilingual-base/blob/main/scripts/gte_embedding.py
2
3from gte_embedding import GTEEmbeddidng
4
5model_name_or_path = 'Alibaba-NLP/gte-multilingual-base'
6model = GTEEmbeddidng(model_name_or_path)
7query = "中国的首都在哪儿"
8
9docs = [
10 "what is the capital of China?",
11 "how to implement quick sort in python?",
12 "北京",
13 "快排算法介绍"
14]
15
16embs = model.encode(docs, return_dense=True,return_sparse=True)
17print('dense_embeddings vecs', embs['dense_embeddings'])
18print('token_weights', embs['token_weights'])
19pairs = [(query, doc) for doc in docs]
20dense_scores = model.compute_scores(pairs, dense_weight=1.0, sparse_weight=0.0)
21sparse_scores = model.compute_scores(pairs, dense_weight=0.0, sparse_weight=1.0)
22hybrid_scores = model.compute_scores(pairs, dense_weight=1.0, sparse_weight=0.3)
23
24print('dense_scores', dense_scores)
25print('sparse_scores', sparse_scores)
26print('hybrid_scores', hybrid_scores)
27
28# dense_scores [0.85302734375, 0.257568359375, 0.76953125, 0.325439453125]
29# sparse_scores [0.0, 0.0, 4.600879669189453, 1.570279598236084]
30# hybrid_scores [0.85302734375, 0.257568359375, 2.1497951507568356, 0.7965233325958252]
31


@inproceedings{zhang2024mgte,
title={mGTE: Generalized Long-Context Text Representation and Reranking Models for Multilingual Text Retrieval},
author={Zhang, Xin and Zhang, Yanzhao and Long, Dingkun and Xie, Wen and Dai, Ziqi and Tang, Jialong and Lin, Huan and Yang, Baosong and Xie, Pengjun and Huang, Fei and others},
booktitle={Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing: Industry Track},
pages={1393--1412},
year={2024}
}