当前位置：首页 > news >正文

Qwen3-Embedding-4B代码实例：调用Embedding API获取向量并本地计算余弦相似度

news 2026/7/3 7:04:40

Qwen3-Embedding-4B代码实例：调用Embedding API获取向量并本地计算余弦相似度

1. 项目概述

Qwen3-Embedding-4B是阿里通义千问系列中的文本嵌入模型，能够将文本转换为高维向量表示。这些向量捕捉了文本的语义信息，使得我们可以通过计算向量间的相似度来实现语义搜索功能，而不仅仅是传统的关键词匹配。

与普通的关键词搜索不同，基于嵌入向量的语义搜索能够理解文本的深层含义。例如，搜索"我想吃点东西"可能会匹配到"苹果是一种很好吃的水果"，尽管这两句话没有任何共同的关键词。

2. 环境准备与安装

2.1 系统要求

Python 3.8或更高版本
CUDA 11.7或更高版本（推荐使用GPU加速）
至少16GB内存（处理大量文本时建议32GB以上）

2.2 安装依赖包

pip install torch transformers sentence-transformers numpy streamlit

2.3 下载模型

Qwen3-Embedding-4B模型可以通过Hugging Face的transformers库直接加载：

from transformers import AutoModel, AutoTokenizer model_name = "Qwen/Qwen3-Embedding-4B" tokenizer = AutoTokenizer.from_pretrained(model_name) model = AutoModel.from_pretrained(model_name).cuda() # 使用GPU加速

3. 核心功能实现

3.1 文本向量化

文本向量化是将文本转换为固定长度向量的过程。Qwen3-Embedding-4B生成的向量维度为4096。

def get_embedding(text): inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True) inputs = {k: v.cuda() for k, v in inputs.items()} # 将输入数据移动到GPU with torch.no_grad(): outputs = model(**inputs) # 获取最后一层隐藏状态并做平均池化 last_hidden_state = outputs.last_hidden_state attention_mask = inputs["attention_mask"] # 扩展attention_mask的维度以匹配hidden_state expanded_mask = attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float() # 对隐藏状态进行加权平均 sum_embeddings = torch.sum(last_hidden_state * expanded_mask, 1) sum_mask = torch.clamp(expanded_mask.sum(1), min=1e-9) embedding = sum_embeddings / sum_mask return embedding.cpu().numpy() # 将结果移回CPU并转换为numpy数组

3.2 余弦相似度计算

余弦相似度是衡量两个向量方向相似程度的指标，值域在[-1,1]之间，值越大表示越相似。

import numpy as np def cosine_similarity(vec1, vec2): # 确保向量是二维的 vec1 = np.array(vec1).reshape(1, -1) vec2 = np.array(vec2).reshape(1, -1) # 计算点积 dot_product = np.dot(vec1, vec2.T) # 计算范数 norm_vec1 = np.linalg.norm(vec1, axis=1) norm_vec2 = np.linalg.norm(vec2, axis=1) # 计算余弦相似度 similarity = dot_product / (norm_vec1 * norm_vec2) return similarity[0][0]

4. 完整语义搜索实现

4.1 构建知识库

首先，我们需要准备一个知识库，包含我们想要搜索的文本内容。

knowledge_base = [ "苹果是一种很好吃的水果", "Python是一种流行的编程语言", "深度学习需要大量的计算资源", "北京是中国的首都", "健康的饮食应该包含多种蔬菜", "定期锻炼对身体有益", "机器学习是人工智能的一个分支", "充足的睡眠对健康很重要" ]

4.2 知识库向量化

为了提高搜索效率，我们可以预先计算知识库中所有文本的向量表示。

# 预计算知识库中所有文本的嵌入向量 knowledge_embeddings = [get_embedding(text) for text in knowledge_base]

4.3 语义搜索函数

def semantic_search(query, knowledge_base, knowledge_embeddings, top_k=5): # 获取查询文本的嵌入向量 query_embedding = get_embedding(query) # 计算查询向量与知识库中所有向量的相似度 similarities = [] for emb in knowledge_embeddings: sim = cosine_similarity(query_embedding, emb) similarities.append(sim) # 将结果按相似度排序 results = sorted(zip(knowledge_base, similarities), key=lambda x: x[1], reverse=True) # 返回top_k个最相似的结果 return results[:top_k]

5. 示例应用

5.1 简单查询示例

query = "我想吃点东西" results = semantic_search(query, knowledge_base, knowledge_embeddings) print(f"查询: {query}") print("最相似的结果:") for text, sim in results: print(f"相似度: {sim:.4f} - 文本: {text}")

输出可能类似于：

查询: 我想吃点东西 最相似的结果: 相似度: 0.7523 - 文本: 苹果是一种很好吃的水果 相似度: 0.6124 - 文本: 健康的饮食应该包含多种蔬菜 相似度: 0.4321 - 文本: 充足的睡眠对健康很重要 相似度: 0.3215 - 文本: 定期锻炼对身体有益 相似度: 0.2103 - 文本: 深度学习需要大量的计算资源

5.2 不同领域的查询示例

queries = [ "编程用什么语言好", "怎么保持身体健康", "人工智能的基础是什么" ] for query in queries: print(f"\n查询: {query}") results = semantic_search(query, knowledge_base, knowledge_embeddings, top_k=3) for text, sim in results: print(f"相似度: {sim:.4f} - 文本: {text}")

6. 性能优化建议

6.1 批量处理

当处理大量文本时，可以使用批量处理来提高效率：

def get_embeddings_batch(texts, batch_size=32): all_embeddings = [] for i in range(0, len(texts), batch_size): batch = texts[i:i+batch_size] inputs = tokenizer(batch, return_tensors="pt", padding=True, truncation=True, max_length=512) inputs = {k: v.cuda() for k, v in inputs.items()} with torch.no_grad(): outputs = model(**inputs) last_hidden_state = outputs.last_hidden_state attention_mask = inputs["attention_mask"] expanded_mask = attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float() sum_embeddings = torch.sum(last_hidden_state * expanded_mask, 1) sum_mask = torch.clamp(expanded_mask.sum(1), min=1e-9) batch_embeddings = (sum_embeddings / sum_mask).cpu().numpy() all_embeddings.extend(batch_embeddings) return all_embeddings

6.2 使用FAISS加速搜索

对于大规模知识库，可以使用Facebook的FAISS库来加速相似度搜索：

import faiss # 构建FAISS索引 dimension = knowledge_embeddings[0].shape[1] index = faiss.IndexFlatIP(dimension) # 使用内积作为相似度度量 # 将向量添加到索引中 embeddings_matrix = np.vstack(knowledge_embeddings) faiss.normalize_L2(embeddings_matrix) # 归一化向量以便使用内积计算余弦相似度 index.add(embeddings_matrix) def faiss_semantic_search(query, knowledge_base, index, top_k=5): query_embedding = get_embedding(query) query_embedding = query_embedding.astype('float32') faiss.normalize_L2(query_embedding) # 搜索最相似的top_k个向量 distances, indices = index.search(query_embedding, top_k) # 返回结果 results = [] for i, idx in enumerate(indices[0]): results.append((knowledge_base[idx], distances[0][i])) return results