Qwen-7B-Embedding

1- 基本使用示范

import torch 
from torch import nn
from modelscope.hub.snapshot_download import snapshot_download
from modelscope import AutoTokenizer 
from os.path import join as p_join


snapshot_download(model_id="sccHyFuture/Qwen-7B-Embedding", cache_dir="./")
tokenizer = AutoTokenizer.from_pretrained("qwen/Qwen-7B", trust_remote_code=True)
ipts = tokenizer('一只猫', return_tensors='pt')['input_ids']
embed_dict = torch.load('./sccHyFuture/Qwen-7B-Embedding/qwen_7b_embed.pth')
vocab_size, embd_dim = embed_dict['weight'].size()
embed = nn.Embedding(vocab_size, embd_dim)
embed(ipts)
embed.load_state_dict(embed_dict)
embed(ipts)

2- chormadb Embedding提换示例

详细见chromadb_embd.py

class QwenEmbeddingFunction(EmbeddingFunction):
    def __init__(self,
            model_name: str = "qwen-7B",
            device: str = "cpu",
            normalize_embeddings: bool = False,
            max_length: int=128
        ):
        self.model_name = model_name
        self.device = device
        self.normalize_embeddings = normalize_embeddings
        self.max_length = max_length
        self.tokenizer = AutoTokenizer.from_pretrained("qwen/Qwen-7B", trust_remote_code=True, pad_token='<|endoftext|>')
        self.emb = self._load_embd().to(self.device)

    @torch.no_grad()
    def __call__(self, input: Documents) -> Embeddings:
        tk = self.tokenizer(input, return_tensors='pt', padding=True, truncation=True, max_length=self.max_length)['input_ids'].to(self.device)
        emb_out = self.emb(tk).detach()
        if self.normalize_embeddings:
            emb_out = torch.nn.functional.normalize(emb_out, p=2, dim=1)
        return cast(
            Embeddings,
            [emb.numpy().flatten().tolist() for emb in emb_out] 
        )

    def _load_embd(self):
        sep = os.path.sep
        embed_name_id = 'sccHyFuture/Qwen-7B-Embedding'
        local_embed_file = f'.{sep}{embed_name_id.replace("/", sep)}{sep}qwen_7b_embed.pth'
        if not os.path.exists(local_embed_file):
            snapshot_download(model_id=embed_name_id, cache_dir=f"./")

        embed_dict = torch.load(local_embed_file, map_location=self.device) 
        vocab_size, embd_dim = embed_dict['weight'].size()
        embed = torch.nn.Embedding(vocab_size, embd_dim)
        embed.load_state_dict(embed_dict)
        return embed