jarvis-models/sample/chroma_client_query.py

import os
import time
import chromadb
from chromadb.config import Settings

from langchain_community.document_loaders.csv_loader import CSVLoader
from langchain_community.document_loaders import PyPDFLoader, DirectoryLoader, TextLoader, UnstructuredHTMLLoader, JSONLoader, Docx2txtLoader, UnstructuredExcelLoader
from langchain_community.vectorstores import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter, CharacterTextSplitter
from langchain_community.embeddings.sentence_transformer import SentenceTransformerEmbeddings, HuggingFaceEmbeddings



def get_all_files(folder_path):
    # 获取文件夹下所有文件和文件夹的名称列表
    files = os.listdir(folder_path)

    # 初始化空列表，用于存储所有文件的绝对路径
    absolute_paths = []

    # 遍历文件和文件夹名称列表
    for file in files:
        # 拼接文件的绝对路径
        absolute_path = os.path.join(folder_path, file)
        # 如果是文件，将其绝对路径添加到列表中
        if os.path.isfile(absolute_path):
            absolute_paths.append(absolute_path)

    return absolute_paths

# start_time = time.time()
# # 加载文档
# folder_path = "./text"
# txt_files = get_all_files(folder_path)
# docs = []
# ids = []
# for txt_file in txt_files:
#     loader = PyPDFLoader(txt_file)

#     documents = loader.load()

#     text_splitter = RecursiveCharacterTextSplitter(chunk_size=512, chunk_overlap=0)

#     docs_txt = text_splitter.split_documents(documents)

#     docs.extend(docs_txt)

#     ids.extend([os.path.basename(txt_file) + str(i) for i in range(len(docs_txt))])
# start_time1 = time.time()
# print(start_time1 - start_time)


# loader = PyPDFLoader("/code/memory/text/大语言模型应用.pdf")
# loader = TextLoader("/code/memory/text/test.txt")
# loader = CSVLoader("/code/memory/text/test1.csv")
# loader = UnstructuredHTMLLoader("/"example_data/fake-content.html"")
# pip install docx2txt
# loader = Docx2txtLoader("/code/memory/text/tesou.docx")
# pip install openpyxl
# loader = UnstructuredExcelLoader("/code/memorinject_prompt = '(用活泼的语气说话回答，回答严格限制50字以内)'
#        inject_prompt = '(回答简练，不要输出重复内容，只讲中文)'
# text_splitter = RecursiveCharacterTextSplitter(chunk_size=512, chunk_overlap=0)
# docs = text_splitter.split_documents(documents)
# print(len(docs))
# ids = ["大语言模型应用"+str(i) for i in range(len(docs))]


# 加载文档和拆分文档
# loader = TextLoader("/home/administrator/Workspace/jarvis-models/sample/RAG_zh.txt")

# documents = loader.load()

# text_splitter = RecursiveCharacterTextSplitter(chunk_size=1024, chunk_overlap=50)

# docs = text_splitter.split_documents(documents)

# print("len(docs)", len(docs))

# ids = ["20240521_store"+str(i) for i in range(len(docs))]


# # 加载embedding模型和chroma server
# embedding_model = SentenceTransformerEmbeddings(model_name='/home/administrator/Workspace/Models/BAAI/bge-large-zh-v1.5', model_kwargs={"device": "cuda"})
# client = chromadb.HttpClient(host='172.16.5.8', port=7000)

# id = "g2e"
# client.delete_collection(id)
# collection_number = client.get_or_create_collection(id).count()
# print("collection_number",collection_number)
# start_time2 = time.time()
# # 插入向量(如果ids已存在，则会更新向量)
# db = Chroma.from_documents(documents=docs, embedding=embedding_model, ids=ids, collection_name=id, client=client)

# # db = Chroma.from_texts(texts=['test by tom'], embedding=embedding_model, ids=["大语言模型应用0"], persist_directory="./data/test1", collection_name="123", metadatas=[{"source": "string"}])

# start_time3 = time.time()
# print("insert time ", start_time3 - start_time2)
# collection_number = client.get_or_create_collection(id).count()
# print("collection_number",collection_number)



# chroma 召回
from chromadb.utils import embedding_functions
embedding_model = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="/home/administrator/Workspace/Models/BAAI/bge-large-zh-v1.5", device = "cuda")
client = chromadb.HttpClient(host='172.16.5.8', port=7000)
collection = client.get_collection("g2e", embedding_function=embedding_model)

print(collection.count())
import time
start_time = time.time()
query = "你知道澳门银河吗"
# query it
results = collection.query(
    query_texts=[query],
    n_results=5,    
)

response = results["documents"]
print("response: ", response)
print("time: ", time.time() - start_time)


# # 结合大模型进行总结
# import requests

# model_name = "Qwen1.5-14B-Chat"
# chat_inputs={
#     "model": model_name,
#     "messages": [
#         {
#             "role": "user",
#             "content": f"问题: {query}。- 根据知识库内的检索结果，以清晰简洁的表达方式回答问题。- 只从检索内容中选取与问题密切相关的信息。- 不要编造答案，如果答案不在经核实的资料中或无法从经核实的资料中得出，请回答“我无法回答您的问题。”检索内容：{response}"
#         }
#     ],
#     # "temperature": 0,
#     # "top_p": user_top_p,
#     # "n": user_n,
#     # "max_tokens": user_max_tokens,
#     # "frequency_penalty": user_frequency_penalty,
#     # "presence_penalty": user_presence_penalty,
#     # "stop": 100
# }

# key ="YOUR_API_KEY"

# header = {
#     'Content-Type': 'application/json',
#     'Authorization': "Bearer " + key
# }
# url = "http://172.16.5.8:23333/v1/chat/completions"

# fastchat_response = requests.post(url, json=chat_inputs, headers=header)
# # print(fastchat_response.json())

# print("\n question: ", query)
# print("\n ",model_name, fastchat_response.json()["choices"][0]["message"]["content"])




# start_time4 = time.time()
# db = Chroma(
#     client=client,
#     collection_name=id,
#     embedding_function=embedding_model,
# )

# 更新文档
# db = db.update_documents(ids, documents)
# 删除文档
# db.delete([ids])
# 删除集合
# db.delete_collection()

# query = "智能体核心思想"
# docs = db.similarity_search(query, k=2)

# print("result: ",docs)
# for doc in docs:
#     print(doc, "\n")

# start_time5 = time.time()
# print("search time ", start_time5 - start_time4)

# docs = db._collection.get(ids=['大语言模型应用0'])

# print(docs)

# docs = db.get(where={"source": "text/大语言模型应用.pdf"})
# docs = db.get()
# print(docs)