Views
No views yet
sentence-transformers or transformers.U.pth, V.pth)1from sentence_transformers import SentenceTransformer
2
3model = SentenceTransformer(
4 "codefuse-ai/ML-Embed-0.6B",
5 device="cuda:0",
6 model_kwargs={"torch_dtype": "bfloat16"}
7)
8
9# Some sample query and documents
10query = "What is ML-Embed used for?"
11documents = [
12 "ML-Embed is a family of multilingual embedding models for retrieval, semantic search, and other NLP tasks.",
13 "ML-Embed is trained to produce text embeddings that work well across many languages.",
14 "ML-Embed 是 CodeFuse AI 开源的多语言嵌入模型。",
15 "ML-Embed — это многоязычная модель эмбеддингов для поиска и семантического сопоставления."
16]
17
18# Encode the query and documents separately. The encode_query method uses the query prompt
19query_embedding = model.encode_query(query)
20document_embeddings = model.encode_document(documents)
21
22print(query_embedding.shape, document_embeddings.shape)
23# (1024,) (4, 1024)
24
25# Compute cosine similarity between the query and documents
26similarity = model.similarity(query_embedding, document_embeddings)
27print(similarity)1from transformers import AutoModel, AutoTokenizer
2import torch
3import torch.nn.functional as F
4
5model_path = "codefuse-ai/ML-Embed-0.6B"
6
7tokenizer = AutoTokenizer.from_pretrained(model_path)
8model = AutoModel.from_pretrained(
9 model_path,
10 torch_dtype=torch.bfloat16,
11 device_map={"": 0}
12)
13
14query = "What is ML-Embed used for?"
15query_prompt = "Instruct: Given a question, retrieve passages that can help answer the question.\nQuery: "
16
17documents = [
18 "ML-Embed is a family of multilingual embedding models for retrieval, semantic search, and other NLP tasks.",
19 "ML-Embed is trained to produce text embeddings that work well across many languages.",
20 "ML-Embed 是 CodeFuse AI 开源的多语言嵌入模型。",
21 "ML-Embed — это многоязычная модель эмбеддингов для поиска и семантического сопоставления."
22]
23
24def encode(sentences):
25 batch_size = len(sentences)
26 tokenized_inputs = tokenizer(sentences, padding=True, return_tensors="pt").to(model.device)
27 last_hidden_state = model(**tokenized_inputs).last_hidden_state
28 eos_positions = tokenized_inputs.attention_mask.sum(dim=1) - 1
29 embeddings = last_hidden_state[torch.arange(batch_size, device=model.device), eos_positions]
30 embeddings = F.normalize(embeddings, p=2, dim=1)
31 return embeddings
32
33# Encode the query and documents
34query_embedding = encode([query_prompt + query])
35document_embeddings = encode(documents)
36
37print(query_embedding.shape, document_embeddings.shape)
38# torch.Size([1, 1024]) torch.Size([4, 1024])
39
40# Compute cosine similarity between the query and documents
41similarity = query_embedding @ document_embeddings.T
42print(similarity)1Instruct: your_instruction
2Query:AutoModel or SentenceTransformer without any code changes (refer to the examples above).num_hidden_layersmax_window_layersconfig.json to a value smaller than the current one.28 to 16 will make transformers load only the first 16 layers and ignore the remaining weights.transformers library.Note: make surenum_hidden_layersandmax_window_layersstay consistent. If you are using transformers v5, you will also need to trucatelayer_typesin the config file according to the new layer count.
U.pthV.pthU.pth and V.pth1import torch
2import torch.nn.functional as F
3from transformers import AutoModel, AutoTokenizer
4from huggingface_hub import hf_hub_download
5
6model_path = "codefuse-ai/ML-Embed-0.6B"
7dtype = torch.bfloat16
8device = "cuda"
9
10# Load tokenizer and model
11tokenizer = AutoTokenizer.from_pretrained(model_path)
12model = AutoModel.from_pretrained(
13 model_path,
14 torch_dtype=dtype,
15 device_map={"": 0}
16)
17model.eval()
18
19# Load factorized embedding weights
20u_path = hf_hub_download(repo_id=model_path, filename="U.pth")
21v_path = hf_hub_download(repo_id=model_path, filename="V.pth")
22
23U = torch.load(u_path, map_location="cpu").to(dtype).to(device)
24V = torch.load(v_path, map_location="cpu").to(dtype).to(device)
25
26# Optional: choose a smaller rank for more compression
27# If rank is None, use the full factorized rank
28rank = None
29
30def encode_with_factorized_embedding(sentences, rank=None):
31 tokenized = tokenizer(sentences, padding=True, return_tensors="pt").to(device)
32 input_ids = tokenized["input_ids"]
33 attention_mask = tokenized["attention_mask"]
34
35 if rank is None:
36 inputs_embeds = (U @ V)[input_ids]
37 else:
38 inputs_embeds = (U[:, :rank] @ V[:rank, :])[input_ids]
39
40 outputs = model(
41 inputs_embeds=inputs_embeds,
42 attention_mask=attention_mask
43 )
44
45 last_hidden_state = outputs.last_hidden_state
46 eos_positions = attention_mask.sum(dim=1) - 1
47 embeddings = last_hidden_state[torch.arange(len(sentences), device=device), eos_positions]
48 embeddings = F.normalize(embeddings, p=2, dim=1)
49 return embeddings
50
51query = "What is ML-Embed used for?"
52query_prompt = "Instruct: Given a question, retrieve passages that can help answer the question.\nQuery: "
53
54documents = [
55 "ML-Embed is a family of multilingual embedding models for retrieval, semantic search, and other NLP tasks.",
56 "ML-Embed is trained to produce text embeddings that work well across many languages.",
57 "ML-Embed 是 CodeFuse AI 开源的多语言嵌入模型。",
58 "ML-Embed — это многоязычная модель эмбеддингов для поиска и семантического сопоставления."
59]
60
61query_embedding = encode_with_factorized_embedding([query_prompt + query], rank=rank)
62document_embeddings = encode_with_factorized_embedding(documents, rank=rank)
63
64similarity = query_embedding @ document_embeddings.T
65print(similarity)d dimensions. This can reduce storage and speed up vector search in downstream systems. The model is trained with a smallest Matryoshka dimension of 8.1embedding = embedding[..., :512]
2embedding = torch.nn.functional.normalize(embedding, p=2, dim=-1)Note: you need to apply normalization after trucation, not the other way around.
truncate_dim=512 to the encode interface.codefuse-ai/F2LLM-v21@misc{zhang2026mlembedinclusiveefficientembeddings,
2 title={ML-Embed: Inclusive and Efficient Embeddings for a Multilingual World},
3 author={Ziyin Zhang and Zihan Liao and Hang Yu and Peng Di and Rui Wang},
4 year={2026},
5 eprint={2605.15081},
6 archivePrefix={arXiv},
7 primaryClass={cs.CL},
8 url={https://arxiv.org/abs/2605.15081},
9}F2LLM-v2: