基于Ollama和Python的文本嵌入与向量检索实践

文本嵌入技术是自然语言处理领域的重要基础，能够将文本转换为高维向量表示，从而支持语义搜索、相似度计算等应用。本文将介绍如何利用Ollama项目和Python实现一个完整的文本嵌入与检索系统。

核心组件介绍

本系统主要依赖两个关键组件：

1. Ollama：一个开源的文本嵌入模型服务，提供多种预训练模型用于生成文本向量表示
2. ChromaDB：一个轻量级的向量数据库，专门用于存储和检索高维向量数据

实现原理

系统工作流程分为两个主要阶段：

1. 文本嵌入存储阶段

首先，我们需要将原始文本转换为向量表示并存储：

async def store_embeddings(ollama_client):
    model_name = 'granite-embedding:278m'
    for i, doc in enumerate(documents):
        response = await ollama_client.embed(model=model_name, input=doc)
        embeddings = response.get("embeddings", [])
        collection.add(
            ids=[str(i)],
            embeddings=embeddings,
            documents=[doc]
        )

这段代码使用Ollama的异步客户端，调用指定的嵌入模型将文本转换为向量，然后将结果存储到ChromaDB中。

2. 查询检索阶段

当用户输入查询文本时，系统会：

1. 将查询文本转换为向量
2. 在向量数据库中查找最相似的文档

async def query_embedding(ollama_client, query_text):
    response = await ollama_client.embed(model=model_name, input=query_text)
    embeddings = response.get("embeddings", [])
    results = collection.query(query_embeddings=embeddings, n_results=1)
    return results['documents'][0][0]

关键技术点

1. 模型选择：系统使用了granite-embedding:278m模型，这是一个278M参数的文本嵌入模型，适合中小规模应用

2. 异步处理：采用Python的asyncio实现异步操作，提高系统吞吐量

3. 错误处理：代码中包含了完善的错误处理机制，确保系统稳定性

4. 向量归一化：在查询处理中对返回的向量进行了格式检查和转换

性能优化建议

1. 批量处理：可以考虑将文档批量发送给嵌入模型，减少网络开销

2. 缓存机制：对频繁查询的内容可以添加缓存层

3. 模型调优：根据具体应用场景选择合适的嵌入模型

4. 索引优化：ChromaDB支持多种索引类型，可以根据数据规模选择合适的索引策略

应用场景

这种基于嵌入向量的检索系统可以应用于：

• 知识库问答
• 文档相似性搜索
• 推荐系统
• 语义缓存

总结

本文展示了一个完整的文本嵌入与检索系统的实现方法。通过结合Ollama的嵌入模型和ChromaDB的向量存储能力，开发者可以快速构建高效的语义搜索应用。系统设计考虑了性能、稳定性和扩展性，可以作为更复杂应用的基础框架。