promptSim/text_analysis/chroma1.py

# from chromadb.config import Settings
from sentence_transformers import SentenceTransformer
import chromadb
from tqdm import tqdm
from langchain.text_splitter import RecursiveCharacterTextSplitter
import uuid

class LangChainChroma:
    def __init__(self,collection_name):
        dirPath="../chromaDB/allField/"
        self.chroma_client = chromadb.PersistentClient(path=dirPath+collection_name)
        # chroma_client=chromadb.Client(Settings(allow_reset=True,persist_directory="../allField/demo/"))
        self.collection=self.chroma_client.get_or_create_collection(name=collection_name,metadata={"hnsw:space": "cosine"})
        model = SentenceTransformer('text_analysis/shibing624/text2vec-base-chinese')
        # model = SentenceTransformer('shibing624/text2vec-base-chinese')
        self.bge_model = model
        self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=500,chunk_overlap=0,separators=["\n\n", "\n", " ", "。"])
        # chroma_client.reset()

    def db_close(self):
        self.chroma_client.clear_system_cache()

    def embedding_fn(self, paragraphs):
        '''文本向量化'''
        doc_vecs = [
            self.bge_model.encode(doc, normalize_embeddings=True).tolist()
            for doc in paragraphs
        ]
        return doc_vecs

    def add_documents(self,documents):
        # embeddings=get_embeddings(documents)
        #向collection中添加文档与向量
        ids = ["id-{}".format(uuid.uuid1()) for i in range(len(documents))]
        self.collection.add(
            embeddings=self.embedding_fn(documents),#每个文档的向量
            documents=documents,
            ids=ids
        )
        # logging.info('当前数据划分{}个块。数据库共有{}个块'.format(len(documents),db_count))
        return ids

    def search(self,queryQ,top_n):
        results=self.collection.query(
            # query_texts=[query],
            query_embeddings=self.embedding_fn([queryQ]),
            n_results=top_n
        )
        return results


# vector_db=LangChainChroma("demo")
# with open("policy_test2.txt", "r", encoding="utf8") as f:
#     for line in tqdm(f):
#         # documents = [Document(page_content=line)]
#         docs = text_splitter.split_text(line)
#         a=vector_db.add_documents(docs)
#         print(a)
#     print("over")


# user_query="鲍炳章同志？"
# results=vector_db.search(user_query,3)
# for para in results['documents'][0]:
#     print(para+'\n')