tinyvector开发者指南：如何为项目添加自定义嵌入生成功能（支持OpenAI/Cohere）
项目地址: https://gitcode.com/gh_mirrors/ti/tinyvector 🚀 终极指南：为tinyvector向量数据库添加自定义嵌入生成功能，轻松支持OpenAI、Cohere等主流模型！本文将带你一步步实现这个强大的功能扩展，让tinyvector从一个简单的向量存储库升级为完整的嵌入生成解决方案。
为什么需要自定义嵌入生成功能？
tinyvector是一个轻量级的最近邻嵌入数据库，基于SQLite和Numpy构建。虽然它目前专注于向量存储和检索，但在实际应用中，用户通常需要将文本转换为嵌入向量才能使用。通过添加自定义嵌入生成功能，你可以：
- 无缝集成：直接使用OpenAI、Cohere等API生成嵌入
- 本地模型支持：集成Hugging Face或Sentence Transformers模型
- 简化工作流程：从文本到检索的一站式解决方案
- 降低成本：减少外部API调用和数据传输
tinyvector项目架构概览
在开始之前，让我们先了解tinyvector的核心结构：
tinyvector/
├── database.py # 核心数据库类，包含索引和查询逻辑
├── types/
│ └── model_db.py # 数据模型定义（Pydantic）
└── server/
├── __main__.py # Flask服务器主文件
└── models/
└── model_response.py # API响应模型
项目的核心是tinyvector/database.py中的DB类，它负责处理所有的数据库操作。当前系统通过insert方法接收预计算的嵌入向量，但没有内置的文本到嵌入转换功能。
步骤1：创建嵌入生成器抽象类
首先，我们需要创建一个通用的嵌入生成器接口。在tinyvector/目录下创建新的文件：
# tinyvector/embedding_generator.py
from abc import ABC, abstractmethod
from typing import List, Optional
import numpy as np
class EmbeddingGenerator(ABC):
"""嵌入生成器抽象基类"""
@abstractmethod
def generate_embedding(self, text: str) -> List[float]:
"""生成单个文本的嵌入向量"""
pass
@abstractmethod
def generate_embeddings(self, texts: List[str]) -> List[List[float]]:
"""批量生成文本的嵌入向量"""
pass
@abstractmethod
def get_dimension(self) -> int:
"""返回嵌入向量的维度"""
pass
这个抽象类定义了所有嵌入生成器必须实现的方法，确保不同模型实现的一致性。
步骤2：实现OpenAI嵌入生成器
接下来，我们实现OpenAI的嵌入生成器。首先安装必要的依赖：
pip install openai numpy
然后创建OpenAI嵌入生成器的具体实现：
# tinyvector/embedding_providers/openai_provider.py
import os
from typing import List, Optional
import numpy as np
from openai import OpenAI
from ..embedding_generator import EmbeddingGenerator
class OpenAIEmbeddingGenerator(EmbeddingGenerator):
"""OpenAI嵌入生成器实现"""
def __init__(
self,
api_key: Optional[str] = None,
model: str = "text-embedding-3-small",
base_url: Optional[str] = None
):
"""
初始化OpenAI嵌入生成器
Args:
api_key: OpenAI API密钥，如果为None则从环境变量读取
model: 使用的嵌入模型名称
base_url: OpenAI API基础URL（用于自定义部署）
"""
self.api_key = api_key or os.getenv("OPENAI_API_KEY")
if not self.api_key:
raise ValueError("OpenAI API密钥未提供且未在环境变量中找到")
self.model = model
self.client = OpenAI(api_key=self.api_key, base_url=base_url)
# 预定义模型维度映射
self._model_dimensions = {
"text-embedding-3-small": 1536,
"text-embedding-3-large": 3072,
"text-embedding-ada-002": 1536
}
def generate_embedding(self, text: str) -> List[float]:
"""生成单个文本的嵌入向量"""
response = self.client.embeddings.create(
model=self.model,
input=text
)
return response.data[0].embedding
def generate_embeddings(self, texts: List[str]) -> List[List[float]]:
"""批量生成文本的嵌入向量"""
response = self.client.embeddings.create(
model=self.model,
input=texts
)
return [item.embedding for item in response.data]
def get_dimension(self) -> int:
"""返回嵌入向量的维度"""
if self.model in self._model_dimensions:
return self._model_dimensions[self.model]
# 如果模型不在映射中，通过API获取维度
try:
test_embedding = self.generate_embedding("test")
return len(test_embedding)
except Exception:
# 默认返回常见维度
return 1536
步骤3：实现Cohere嵌入生成器
同样地，我们实现Cohere的嵌入生成器：
# tinyvector/embedding_providers/cohere_provider.py
import os
from typing import List, Optional
import numpy as np
import cohere
from ..embedding_generator import EmbeddingGenerator
class CohereEmbeddingGenerator(EmbeddingGenerator):
"""Cohere嵌入生成器实现"""
def __init__(
self,
api_key: Optional[str] = None,
model: str = "embed-english-v3.0",
base_url: Optional[str] = None
):
"""
初始化Cohere嵌入生成器
Args:
api_key: Cohere API密钥
model: 使用的嵌入模型名称
base_url: Cohere API基础URL
"""
self.api_key = api_key or os.getenv("COHERE_API_KEY")
if not self.api_key:
raise ValueError("Cohere API密钥未提供且未在环境变量中找到")
self.model = model
self.co = cohere.Client(self.api_key, base_url=base_url)
# Cohere模型维度映射
self._model_dimensions = {
"embed-english-v3.0": 1024,
"embed-english-light-v3.0": 384,
"embed-multilingual-v3.0": 1024,
"embed-multilingual-light-v3.0": 384
}
def generate_embedding(self, text: str) -> List[float]:
"""生成单个文本的嵌入向量"""
response = self.co.embed(
texts=[text],
model=self.model,
input_type="search_document"
)
return response.embeddings[0]
def generate_embeddings(self, texts: List[str]) -> List[List[float]]:
"""批量生成文本的嵌入向量"""
response = self.co.embed(
texts=texts,
model=self.model,
input_type="search_document"
)
return response.embeddings
def get_dimension(self) -> int:
"""返回嵌入向量的维度"""
if self.model in self._model_dimensions:
return self._model_dimensions[self.model]
# 通过API获取维度
try:
test_embedding = self.generate_embedding("test")
return len(test_embedding)
except Exception:
return 1024 # 默认维度
步骤4：实现本地模型嵌入生成器
对于想要在本地运行的开发者，我们可以添加Hugging Face模型支持：
# tinyvector/embedding_providers/huggingface_provider.py
from typing import List, Optional
import numpy as np
from transformers import AutoTokenizer, AutoModel
import torch
from ..embedding_generator import EmbeddingGenerator
class HuggingFaceEmbeddingGenerator(EmbeddingGenerator):
"""Hugging Face本地模型嵌入生成器"""
def __init__(
self,
model_name: str = "sentence-transformers/all-MiniLM-L6-v2",
device: Optional[str] = None
):
"""
初始化Hugging Face嵌入生成器
Args:
model_name: Hugging Face模型名称
device: 运行设备（cpu/cuda）
"""
self.model_name = model_name
self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
# 加载tokenizer和模型
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.model = AutoModel.from_pretrained(model_name).to(self.device)
self.model.eval()
# 获取模型维度
with torch.no_grad():
test_input = self.tokenizer("test", return_tensors="pt", padding=True, truncation=True).to(self.device)
test_output = self.model(**test_input)
self._dimension = test_output.last_hidden_state.size(-1)
def generate_embedding(self, text: str) -> List[float]:
"""生成单个文本的嵌入向量"""
return self.generate_embeddings([text])[0]
def generate_embeddings(self, texts: List[str]) -> List[List[float]]:
"""批量生成文本的嵌入向量"""
# 编码文本
encoded_input = self.tokenizer(
texts,
padding=True,
truncation=True,
return_tensors="pt"
).to(self.device)
# 生成嵌入
with torch.no_grad():
model_output = self.model(**encoded_input)
# 使用平均池化获取句子嵌入
attention_mask = encoded_input['attention_mask']
token_embeddings = model_output.last_hidden_state
# 应用注意力掩码
input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1)
sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
embeddings = sum_embeddings / sum_mask
# 转换为列表并归一化
embeddings = embeddings.cpu().numpy()
embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
return embeddings.tolist()
def get_dimension(self) -> int:
"""返回嵌入向量的维度"""
return self._dimension
步骤5：创建嵌入生成器工厂
为了方便使用，我们创建一个工厂类来管理不同的嵌入生成器：
# tinyvector/embedding_factory.py
from typing import Dict, Optional, Union
from .embedding_generator import EmbeddingGenerator
from .embedding_providers.openai_provider import OpenAIEmbeddingGenerator
from .embedding_providers.cohere_provider import CohereEmbeddingGenerator
from .embedding_providers.huggingface_provider import HuggingFaceEmbeddingGenerator
class EmbeddingFactory:
"""嵌入生成器工厂类"""
PROVIDERS = {
"openai": OpenAIEmbeddingGenerator,
"cohere": CohereEmbeddingGenerator,
"huggingface": HuggingFaceEmbeddingGenerator
}
@classmethod
def create_generator(
cls,
provider: str,
**kwargs
) -> EmbeddingGenerator:
"""
创建嵌入生成器实例
Args:
provider: 提供者名称（openai, cohere, huggingface）
**kwargs: 传递给具体生成器的参数
Returns:
EmbeddingGenerator实例
"""
provider = provider.lower()
if provider not in cls.PROVIDERS:
raise ValueError(
f"不支持的嵌入提供者: {provider}. "
f"支持的提供者: {list(cls.PROVIDERS.keys())}"
)
generator_class = cls.PROVIDERS[provider]
return generator_class(**kwargs)
@classmethod
def get_available_providers(cls) -> list:
"""获取所有可用的嵌入提供者"""
return list(cls.PROVIDERS.keys())
步骤6：扩展数据库类支持嵌入生成
现在我们需要扩展tinyvector的数据库类，使其支持嵌入生成功能：
# 在tinyvector/database.py中添加
from typing import Optional, List
from .embedding_factory import EmbeddingFactory
class DB:
# ... 现有代码 ...
def __init__(self, path, debug=False):
# ... 现有初始化代码 ...
self.embedding_generator = None
def set_embedding_generator(self, provider: str, **kwargs):
"""
设置嵌入生成器
Args:
provider: 嵌入提供者名称
**kwargs: 嵌入生成器参数
"""
self.embedding_generator = EmbeddingFactory.create_generator(provider, **kwargs)
def insert_with_text(
self,
table_name: str,
text: str,
content: Optional[str] = None,
id: Optional[str] = None,
defer_index_update: bool = False
):
"""
插入文本并自动生成嵌入
Args:
table_name: 表名
text: 要插入的文本
content: 可选的内容字段
id: 可选的ID，如果为None则自动生成
defer_index_update: 是否延迟索引更新
"""
if self.embedding_generator is None:
raise ValueError("请先设置嵌入生成器，使用set_embedding_generator()方法")
# 生成嵌入
embedding = self.embedding_generator.generate_embedding(text)
# 插入到数据库
return self.insert(
table_name=table_name,
id=id,
embedding=embedding,
content=content or text,
defer_index_update=defer_index_update
)
def batch_insert_with_texts(
self,
table_name: str,
texts: List[str],
contents: Optional[List[str]] = None,
ids: Optional[List[str]] = None,
defer_index_update: bool = False
):
"""
批量插入文本并自动生成嵌入
Args:
table_name: 表名
texts: 文本列表
contents: 可选的内容列表
ids: 可选的ID列表
defer_index_update: 是否延迟索引更新
"""
if self.embedding_generator is None:
raise ValueError("请先设置嵌入生成器，使用set_embedding_generator()方法")
# 批量生成嵌入
embeddings = self.embedding_generator.generate_embeddings(texts)
# 批量插入
results = []
for i, (text, embedding) in enumerate(zip(texts, embeddings)):
content = contents[i] if contents else text
id = ids[i] if ids else None
result = self.insert(
table_name=table_name,
id=id,
embedding=embedding,
content=content,
defer_index_update=defer_index_update
)
results.append(result)
return results
def query_with_text(
self,
table_name: str,
query_text: str,
k: int = 10
):
"""
使用文本查询
Args:
table_name: 表名
query_text: 查询文本
k: 返回的最相似结果数量
"""
if self.embedding_generator is None:
raise ValueError("请先设置嵌入生成器，使用set_embedding_generator()方法")
# 生成查询嵌入
query_embedding = self.embedding_generator.generate_embedding(query_text)
# 执行查询
return self.query(
table_name=table_name,
query=query_embedding,
k=k
)
步骤7：扩展服务器API
最后，我们需要扩展Flask服务器以支持新的文本插入和查询端点：
# 在server/__main__.py中添加
from pydantic import BaseModel
from typing import Optional, List
from tinyvector.embedding_factory import EmbeddingFactory
class TextInsertionBody(BaseModel):
"""文本插入请求体"""
table_name: str
text: str
content: Optional[str] = None
id: Optional[str] = None
defer_index_update: bool = False
provider: str = "openai"
provider_config: Optional[dict] = None
class BatchTextInsertionBody(BaseModel):
"""批量文本插入请求体"""
table_name: str
texts: List[str]
contents: Optional[List[str]] = None
ids: Optional[List[str]] = None
defer_index_update: bool = False
provider: str = "openai"
provider_config: Optional[dict] = None
class TextQueryObject(BaseModel):
"""文本查询请求体"""
table_name: str
query_text: str
k: int = 10
provider: str = "openai"
provider_config: Optional[dict] = None
# 添加新的API端点
@app.route("/insert_with_text", methods=["POST"])
@spec.validate(
body=TextInsertionBody,
resp=Response(HTTP_200=SuccessMessage, HTTP_400=ErrorMessage),
tags=["embedding"]
)
def insert_with_text():
"""插入文本并自动生成嵌入"""
try:
data = request.get_json()
body = TextInsertionBody(**data)
# 获取或创建嵌入生成器
db = get_db()
if not hasattr(db, 'embedding_generator') or db.embedding_generator is None:
# 动态创建嵌入生成器
generator = EmbeddingFactory.create_generator(
body.provider,
**(body.provider_config or {})
)
db.embedding_generator = generator
# 插入文本
db.insert_with_text(
table_name=body.table_name,
text=body.text,
content=body.content,
id=body.id,
defer_index_update=body.defer_index_update
)
return jsonify({"status": "success"})
except Exception as e:
logging.error(f"Error inserting text: {e}")
return jsonify({"error": str(e)}), 400
@app.route("/query_with_text", methods=["POST"])
@spec.validate(
body=TextQueryObject,
resp=Response(HTTP_200=TableQueryResult, HTTP_400=ErrorMessage),
tags=["embedding"]
)
def query_with_text():
"""使用文本查询"""
try:
data = request.get_json()
body = TextQueryObject(**data)
# 获取或创建嵌入生成器
db = get_db()
if not hasattr(db, 'embedding_generator') or db.embedding_generator is None:
# 动态创建嵌入生成器
generator = EmbeddingFactory.create_generator(
body.provider,
**(body.provider_config or {})
)
db.embedding_generator = generator
# 执行查询
results = db.query_with_text(
table_name=body.table_name,
query_text=body.query_text,
k=body.k
)
return jsonify({"items": results})
except Exception as e:
logging.error(f"Error querying with text: {e}")
return jsonify({"error": str(e)}), 400
步骤8：更新依赖和配置
更新requirements.txt文件以包含新的依赖：
# 添加以下依赖
openai>=1.0.0
cohere>=5.0.0
transformers>=4.30.0
torch>=2.0.0
sentence-transformers>=2.2.0
使用示例
现在你可以轻松使用新的嵌入生成功能了：
from tinyvector import DB
from tinyvector.embedding_factory import EmbeddingFactory
# 创建数据库实例
db = DB("my_database.db")
# 设置OpenAI嵌入生成器
db.set_embedding_generator(
provider="openai",
api_key="your-openai-api-key",
model="text-embedding-3-small"
)
# 创建表
db.create_table("documents", dimension=1536, use_uuid=True)
# 插入文本（自动生成嵌入）
db.insert_with_text(
table_name="documents",
text="机器学习是人工智能的一个分支",
content="机器学习简介"
)
# 批量插入
texts = [
"深度学习是机器学习的一个子集",
"自然语言处理处理人类语言",
"计算机视觉让机器看懂世界"
]
db.batch_insert_with_texts("documents", texts)
# 使用文本查询
results = db.query_with_text(
table_name="documents",
query_text="什么是人工智能？",
k=5
)
# 打印结果
for result in results:
print(f"内容: {result['content']}, 相似度: {result['score']:.4f}")
高级功能：自定义嵌入生成器
如果你有自定义的嵌入模型，可以轻松扩展：
from tinyvector.embedding_generator import EmbeddingGenerator
class CustomEmbeddingGenerator(EmbeddingGenerator):
"""自定义嵌入生成器示例"""
def __init__(self, custom_param: str):
self.custom_param = custom_param
# 初始化你的模型
def generate_embedding(self, text: str) -> List[float]:
# 实现你的嵌入生成逻辑
return [0.1, 0.2, 0.3] # 示例
def generate_embeddings(self, texts: List[str]) -> List[List[float]]:
return [self.generate_embedding(text) for text in texts]
def get_dimension(self) -> int:
return 768 # 你的模型维度
# 注册自定义生成器
EmbeddingFactory.PROVIDERS["custom"] = CustomEmbeddingGenerator
性能优化建议
- 批量处理：使用
batch_insert_with_texts进行批量插入，减少API调用 - 缓存嵌入：为常用文本实现嵌入缓存机制
- 异步处理：对于大量数据，考虑使用异步API调用
- 本地模型：对于高频使用场景，考虑使用本地Hugging Face模型
测试你的实现
创建测试文件确保一切正常工作：
# test_embeddings.py
import unittest
from tinyvector import DB
from tinyvector.embedding_factory import EmbeddingFactory
class TestEmbeddingIntegration(unittest.TestCase):
def test_openai_integration(self):
"""测试OpenAI集成"""
db = DB(":memory:")
# 测试创建生成器
generator = EmbeddingFactory.create_generator(
"openai",
api_key="test-key" # 使用测试密钥
)
self.assertIsNotNone(generator)
self.assertEqual(generator.get_dimension(), 1536)
def test_database_integration(self):
"""测试数据库集成"""
db = DB(":memory:")
# 设置模拟生成器
class MockGenerator:
def generate_embedding(self, text):
return [0.1] * 1536
def generate_embeddings(self, texts):
return [[0.1] * 1536 for _ in texts]
def get_dimension(self):
return 1536
db.embedding_generator = MockGenerator()
# 创建表
db.create_table("test", dimension=1536)
# 测试文本插入
result = db.insert_with_text("test", "测试文本")
self.assertIsNotNone(result)
if __name__ == "__main__":
unittest.main()
总结
通过本文的完整指南，你已经成功为tinyvector添加了强大的自定义嵌入生成功能。这个扩展让你能够：
✅ 支持多种嵌入提供商：OpenAI、Cohere、Hugging Face等
✅ 简化工作流程：从文本直接到向量存储和检索
✅ 易于扩展：可以轻松添加新的嵌入模型
✅ 保持轻量级：不改变tinyvector的核心架构
现在你的tinyvector项目已经具备了完整的嵌入生成能力，可以更好地服务于各种AI应用场景。无论是构建智能搜索引擎、文档检索系统，还是创建个性化推荐引擎，这个扩展都能为你提供强大的支持。
记住，tinyvector的魅力在于它的简洁和可扩展性。通过这个自定义嵌入生成功能，你不仅增强了项目的实用性，还为社区贡献了一个有价值的功能扩展。继续探索和优化，让tinyvector成为你AI项目中的得力助手！ 🚀
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
