26 lines
886 B
Python
26 lines
886 B
Python
"""向量存储服务"""
|
||
from langchain_chroma import Chroma
|
||
import config_data as config
|
||
|
||
class VectorStoreService(object):
|
||
def __init__(self, embedding):
|
||
"""
|
||
:param embedding: 嵌入模型的嵌入
|
||
"""
|
||
|
||
self.embedding = embedding
|
||
self.vector_store = Chroma(
|
||
collection_name=config.collection_name,
|
||
embedding_function=self.embedding,
|
||
persist_directory=config.persist_directory,
|
||
)
|
||
|
||
def get_retriever(self):
|
||
return self.vector_store.as_retriever(search_kwargs={"k": config.similarity_threshold})
|
||
|
||
if __name__ == '__main__':
|
||
from langchain_community.embeddings import DashScopeEmbeddings
|
||
embedding = DashScopeEmbeddings(model = "text-embedding-v4")
|
||
retriver = VectorStoreService(embedding).get_retriever()
|
||
doc = retriver.invoke("我的体重180斤，尺码推荐？")
|
||
|