1.创建chroma客户端并且获取要查询的集合


#向量库
chroma = chromadb.HttpClient(host="localhost",port=8000)
collection = chroma.get_collection(name="my_collection")

集合类似于关系型数据库中的表。get_collection这个代码是用来获取集合的,这里写入之前嵌入时的集合名,也可以用get_or_create_collection方法获取集合(如果存在),否则创建集合。

2.获取集合中的前10条数据

collection = chroma.get_collection(name="my_collection")
data = collection.peek()  # 获取集合中的前10条数据
print(data)

3.引入ollama并且获取配置的开源大模型和嵌入模型

import ollama,chromadb
from dataloader import getconfig
#嵌入模型与大模型
embedmodel = getconfig()["embedmodel"]
llmmodel = getconfig()['llmmodel']

4.这里使用一个交互式的查询程序,借助向量库查询到向量后将其作为上下文添加到大模型的prompt中,这是rag大模型应用的经典查询方法


while True:
    query = input("enter your query:")
    if query.lower() == 'quit':
        break
    else:
        #生成查询向量
        queryembed = ollama.embeddings(model=embedmodel,prompt=query)['embedding']
        #用查询向量检索上下文
        # print(queryembed)
        print(collection.get(
            include=["documents"]
          )
        )
        relevantdocs = collection.query(query_embeddings=[queryembed],n_results=3)['documents'][0]
       
        docs = "\n\n".join(relevantdocs)
        #生成prompt
        modelquery= f"""
        请基于以下的上下文回答问题,如果上下文中不包含足够的回答问题的信息,请回答'我暂时无法回答该问题',不要编造
        上下文:
        ====
        {docs}
        ====
        我的问题是:{query}
        """
        print(modelquery)
        #交给大模型进行生成
        stream = ollama.generate(model=llmmodel,prompt=modelquery,stream=True)

        #流式输出生成的结果
        for chunk in stream:
            if chunk["response"]:
                print(chunk['response'],end='',flush=True)


#

Logo

Agent 垂直技术社区,欢迎活跃、内容共建。

更多推荐