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 streamlit2.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_embeddings6.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 results7. 总结
本文介绍了如何使用Qwen3-Embedding-4B模型实现语义搜索功能。关键步骤包括：
- 使用transformers库加载Qwen3-Embedding-4B模型
- 实现文本向量化函数，将文本转换为高维向量表示
- 使用余弦相似度计算向量间的相似程度
- 构建完整的语义搜索流程，包括知识库准备和查询处理
- 提供性能优化建议，如批量处理和FAISS索引
这种基于嵌入向量的语义搜索方法能够理解文本的深层含义，而不仅仅是表面的关键词匹配，适用于各种需要理解自然语言语义的应用场景。
获取更多AI镜像
想探索更多AI镜像和应用场景？访问 CSDN星图镜像广场，提供丰富的预置镜像，覆盖大模型推理、图像生成、视频生成、模型微调等多个领域，支持一键部署。
