114 lines
3.9 KiB
Python
114 lines
3.9 KiB
Python
# Read Chunking Embedding and save it to Vector Database
|
|
import os
|
|
|
|
from pathlib import Path
|
|
from llama_index.readers.file import PDFReader
|
|
|
|
from models.embedding import EmbeddingModel
|
|
from utils.logging_config import setup_logger
|
|
|
|
from pymilvus import MilvusClient
|
|
|
|
|
|
logger = setup_logger("PreRetrieval")
|
|
|
|
def pdfreader(file_path):
|
|
"""读取PDF文件并返回text文本"""
|
|
assert os.path.exists(file_path), "File not found"
|
|
assert file_path.endswith(".pdf"), "File format not supported"
|
|
|
|
doc = PDFReader().load_data(file=Path(file_path))
|
|
|
|
# 简单的拼接起来之后返回纯文本
|
|
text = "\n\n".join([d.get_content() for d in doc])
|
|
return text
|
|
|
|
def plainreader(file_path):
|
|
"""读取普通文本文件并返回text文本"""
|
|
assert os.path.exists(file_path), "File not found"
|
|
|
|
with open(file_path, "r") as f:
|
|
text = f.read()
|
|
return text
|
|
|
|
|
|
class PreRetrieval:
|
|
|
|
def __init__(self, config):
|
|
self.config = config
|
|
self._init_config(config)
|
|
|
|
self.embed_model = EmbeddingModel(config)
|
|
self.client = MilvusClient(config.milvus_local_path)
|
|
|
|
def _init_config(self, config):
|
|
self.vector_dim = 1024 # 暂时不知道这个和 embedding model 的 embedding 大小有什么关系
|
|
self.default_query_limit = 2
|
|
self.default_collection_name = "default"
|
|
|
|
def add_file(self, file, collection_name=None):
|
|
"""添加文件到数据库"""
|
|
collection_name = collection_name or self.default_collection_name
|
|
text = self.read_text(file)
|
|
chunks = self.chunking(text)
|
|
|
|
self.add_documents(chunks, collection_name)
|
|
|
|
def add_documents(self, docs, collection_name):
|
|
"""添加已经分块之后的文本"""
|
|
vectors = self.embed_model.encode(docs)
|
|
|
|
data = [
|
|
{"id": i, "vector": vectors[i], "text": docs[i], "subject": "history"}
|
|
for i in range(len(vectors))
|
|
]
|
|
|
|
# for testing, we drop the collection if it already exists
|
|
# if self.client.has_collection(collection_name=collection_name):
|
|
# self.client.drop_collection(collection_name=collection_name)
|
|
|
|
self.client.create_collection(
|
|
collection_name=collection_name,
|
|
dimension=self.vector_dim, # The vectors we will use in this demo has 768 dimensions
|
|
)
|
|
|
|
res = self.client.insert(collection_name=collection_name, data=data)
|
|
return res
|
|
|
|
def search(self, query, collection_name=None, limit=None):
|
|
collection_name = collection_name or self.default_collection_name
|
|
limit = limit or self.default_query_limit
|
|
|
|
query_vectors = self.embed_model.encode_queries([query])
|
|
|
|
res = self.client.search(
|
|
collection_name=collection_name, # target collection
|
|
data=query_vectors, # query vectors
|
|
limit=limit, # number of returned entities
|
|
output_fields=["text", "subject"], # specifies fields to be returned
|
|
)
|
|
|
|
return res
|
|
|
|
def read_text(self, file):
|
|
support_format = [".pdf", ".txt", "*.md"]
|
|
assert os.path.exists(file), "File not found"
|
|
logger.info(f"Try to read file {file}")
|
|
if os.path.isfile(file):
|
|
if file.endswith(".pdf"):
|
|
return pdfreader(file)
|
|
elif file.endswith(".txt") or file.endswith(".md"):
|
|
return plainreader(file)
|
|
else:
|
|
logger.error(f"File format not supported, only support {support_format}")
|
|
raise Exception(f"File format not supported, only support {support_format}")
|
|
else:
|
|
logger.error(f"Directory not supported now!")
|
|
raise NotImplementedError("Directory not supported now!")
|
|
|
|
def chunking(self, text, chunk_size=1024):
|
|
"""将文本切分成固定大小的块"""
|
|
chunks = []
|
|
for i in range(0, len(text), chunk_size):
|
|
chunks.append(text[i:i + chunk_size])
|
|
return chunks |