Youtu-Embedding 正式开源：腾讯优图推出高性能通用文本嵌入模型

pip install --upgrade tencentcloud-sdk-python

# 拉取 Youtu-Embedding 项目
git clone https://github.com/TencentCloudADP/youtu-embedding.git
 
# 检查 Python 版本
python --version
 
# （如使用 pyenv，可设置 Python 3.10）
pyenv local 3.10.14
 
# 创建并激活虚拟环境
python -m venv youtu-env
source youtu-env/bin/activate

pip install -U pip
pip install "transformers==4.51.3" torch numpy scipy scikit-learn huggingface_hub

git clone https://huggingface.co/tencent/Youtu-Embedding

import torch
import numpy as np
from transformers import AutoModel, AutoTokenizer


class LLMEmbeddingModel():

    def __init__(self, 
                model_name_or_path, 
                batch_size=128, 
                max_length=1024, 
                gpu_id=0):
        """Local embedding model with automatic device selection"""
        self.model = AutoModel.from_pretrained(model_name_or_path, trust_remote_code=True)
        self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, padding_side="right", trust_remote_code=True)
       
        # Device selection: CUDA -> MPS -> CPU
        if torch.cuda.is_available():
            self.device = torch.device(f"cuda:{gpu_id}")
        elif torch.backends.mps.is_available():
            self.device = torch.device("mps")
        else:
            self.device = torch.device("cpu")

        self.model.to(self.device).eval()
        
        self.max_length = max_length
        self.batch_size = batch_size
        
        query_instruction = "Given a search query, retrieve passages that answer the question"
        if query_instruction:
            self.query_instruction = f"Instruction: {query_instruction} \nQuery:"
        else:
            self.query_instruction = "Query:"
        
        self.doc_instruction = ""
        print(f"Model loaded: {model_name_or_path}")
        print(f"Device: {self.device}")
    
    def mean_pooling(self, hidden_state, attention_mask):
        s = torch.sum(hidden_state * attention_mask.unsqueeze(-1).float(), dim=1)
        d = attention_mask.sum(dim=1, keepdim=True).float()
        embedding = s / d
        return embedding
    
    @torch.no_grad()
    def encode(self, sentences_batch, instruction):
        inputs = self.tokenizer(
            sentences_batch,
            padding=True,
            truncation=True,
            return_tensors="pt",
            max_length=self.max_length,
            add_special_tokens=True,
        )
        # Move inputs to device
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        
        with torch.no_grad():
            outputs = self.model(**inputs)
            last_hidden_state = outputs[0]
            
            instruction_tokens = self.tokenizer(
                instruction,
                padding=False,
                truncation=True,
                max_length=self.max_length,
                add_special_tokens=True,
            )["input_ids"]
            if len(np.shape(np.array(instruction_tokens))) == 1:
                inputs["attention_mask"][:, :len(instruction_tokens)] = 0
            else:
                instruction_length = [len(item) for item in instruction_tokens]
                assert len(instruction) == len(sentences_batch)
                for idx in range(len(instruction_length)):
                    inputs["attention_mask"][idx, :instruction_length[idx]] = 0
            
            embeddings = self.mean_pooling(last_hidden_state, inputs["attention_mask"])
            embeddings = torch.nn.functional.normalize(embeddings, dim=-1)
        return embeddings
    
    def encode_queries(self, queries):
        queries = queries if isinstance(queries, list) else [queries]
        queries = [f"{self.query_instruction}{query}" for query in queries]
        return self.encode(queries, self.query_instruction)
    
    def encode_passages(self, passages):
        passages = passages if isinstance(passages, list) else [passages]
        passages = [f"{self.doc_instruction}{passage}" for passage in passages]
        return self.encode(passages, self.doc_instruction)
    
    def compute_similarity_for_vectors(self, q_reps, p_reps):
        if len(p_reps.size()) == 2:
            return torch.matmul(q_reps, p_reps.transpose(0, 1))
        return torch.matmul(q_reps, p_reps.transpose(-2, -1))
    
    def compute_similarity(self, queries, passages):
        q_reps = self.encode_queries(queries)
        p_reps = self.encode_passages(passages)
        scores = self.compute_similarity_for_vectors(q_reps, p_reps)
        scores = scores.detach().cpu().tolist()
        return scores
    
    def display_results(self, query, passages, scores):
        """Display similarity results in a simple format"""
        print(f"\nQuery: {query}")
        print("-" * 50)

        # Sort by similarity score (highest first)
        ranked_results = list(zip(passages, scores[0]))
        ranked_results.sort(key=lambda x: x[1], reverse=True)

        for i, (passage, score) in enumerate(ranked_results, 1):
            print(f"{i}. Score: {score:.4f} - {passage}")

        print("-" * 50)


def main():
    queries = ["What's the weather like?"]
    passages = [
        'The weather is lovely today.',
        "It's so sunny outside!",
        'He drove to the stadium.'
    ]
    
    model_name_or_path = "./Youtu-Embedding"
    model = LLMEmbeddingModel(model_name_or_path)
    scores = model.compute_similarity(queries, passages)

    # Display results with enhanced formatting
    model.display_results(queries[0], passages, scores)

    # Also show raw scores for reference
    print(f"\nRaw scores: {scores}")


if __name__ == "__main__":
    main()
 
 

python test_transformers_local.py

Model loaded: ./Youtu-Embedding
Device: mps
 
Query: What's the weather like?
--------------------------------------------------
1. Score: 0.4465 - The weather is lovely today.
2. Score: 0.3124 - It's so sunny outside!
3. Score: 0.0304 - He drove to the stadium.
--------------------------------------------------
 
Raw scores: [[0.44651979207992554, 0.31240469217300415, 0.030404280871152878]]