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 79971# 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


@misc{zhang2024mgte,
title={mGTE: Generalized Long-Context Text Representation and Reranking Models for Multilingual Text Retrieval},
author={Xin Zhang and Yanzhao Zhang and Dingkun Long and Wen Xie and Ziqi Dai and Jialong Tang and Huan Lin and Baosong Yang and Pengjun Xie and Fei Huang and Meishan Zhang and Wenjie Li and Min Zhang},
year={2024},
eprint={2407.19669},
archivePrefix={arXiv},
primaryClass={cs.CL},
url={https://arxiv.org/abs/2407.19669},
}