1. 添加 chunk 信息
2. 前端的chunk展示优化,默认使用 start_char_idx,如果没有则使用 id 3. 修改默认overlap为100token
This commit is contained in:
parent
568d859826
commit
87807953b5
@ -1,7 +1,8 @@
|
||||
import os
|
||||
import json
|
||||
import time
|
||||
from src.utils import hashstr, logger, is_text_pdf
|
||||
from src.utils import hashstr, logger
|
||||
from src.core.indexing import chunk
|
||||
from src.models.embedding import get_embedding_model
|
||||
|
||||
|
||||
@ -140,7 +141,6 @@ class DataBaseManager:
|
||||
# 先保存一次数据库状态,确保waiting状态被记录
|
||||
self._save_databases()
|
||||
|
||||
from src.core.indexing import chunk
|
||||
for new_file in new_files:
|
||||
file_id = new_file["file_id"]
|
||||
idx = self.get_idx_by_fileid(db, file_id)
|
||||
@ -156,9 +156,10 @@ class DataBaseManager:
|
||||
nodes = chunk(new_file["path"], params=params)
|
||||
|
||||
self.knowledge_base.add_documents(
|
||||
docs=[node.text for node in nodes],
|
||||
file_id=file_id,
|
||||
collection_name=db.metaname,
|
||||
file_id=file_id)
|
||||
docs=[node.text for node in nodes],
|
||||
chunk_infos=[node.dict() for node in nodes])
|
||||
|
||||
idx = self.get_idx_by_fileid(db, file_id)
|
||||
db.files[idx]["status"] = "done"
|
||||
@ -219,20 +220,16 @@ class DataBaseManager:
|
||||
lines = self.knowledge_base.client.query(
|
||||
collection_name=db.metaname,
|
||||
filter=f"file_id == '{file_id}'",
|
||||
output_fields=["id", "text", "file_id", "hash"]
|
||||
output_fields=None
|
||||
)
|
||||
# 删除 vector 字段
|
||||
for line in lines:
|
||||
line.pop("vector")
|
||||
|
||||
lines.sort(key=lambda x: x.get("start_char_idx", 0))
|
||||
logger.debug(f"lines[0]: {lines[0]}")
|
||||
return {"lines": lines}
|
||||
|
||||
def chunking(self, text, params=None):
|
||||
chunk_method = params.get("chunk_method", "fixed")
|
||||
chunk_size = params.get("chunk_size", 500)
|
||||
|
||||
"""将文本切分成固定大小的块"""
|
||||
chunks = []
|
||||
for i in range(0, len(text), chunk_size):
|
||||
chunks.append(text[i:i + chunk_size])
|
||||
return chunks
|
||||
|
||||
def delete_database(self, db_id):
|
||||
db = self.get_kb_by_id(db_id)
|
||||
if db is None:
|
||||
|
||||
@ -8,9 +8,21 @@ from llama_index.readers.file import FlatReader, DocxReader
|
||||
from src.utils import hashstr
|
||||
|
||||
def chunk(text_or_path, params=None):
|
||||
"""
|
||||
将文本或文件切分成固定大小的块
|
||||
|
||||
Args:
|
||||
text_or_path: 文本或文件路径
|
||||
params: 参数
|
||||
chunk_size: 块大小
|
||||
chunk_overlap: 块重叠大小
|
||||
use_parser: 是否使用文件解析器
|
||||
Returns:
|
||||
nodes: 节点列表
|
||||
"""
|
||||
params = params or {}
|
||||
chunk_size = int(params.get("chunk_size", 500))
|
||||
chunk_overlap = int(params.get("chunk_overlap", 20))
|
||||
chunk_overlap = int(params.get("chunk_overlap", 100))
|
||||
splitter = SentenceSplitter(
|
||||
chunk_size=chunk_size,
|
||||
chunk_overlap=chunk_overlap,
|
||||
|
||||
@ -59,7 +59,7 @@ class KnowledgeBase:
|
||||
dimension= dimension, # The vectors we will use in this demo has 768 dimensions
|
||||
)
|
||||
|
||||
def add_documents(self, docs, collection_name, **kwargs):
|
||||
def add_documents(self, docs, collection_name, chunk_infos=None, **kwargs):
|
||||
"""添加已经分块之后的文本"""
|
||||
# 检查 collection 是否存在
|
||||
import random
|
||||
@ -67,6 +67,8 @@ class KnowledgeBase:
|
||||
logger.error(f"Collection {collection_name} not found, create it")
|
||||
# self.add_collection(collection_name)
|
||||
|
||||
chunk_infos = chunk_infos or [{}] * len(docs)
|
||||
|
||||
vectors = self.embed_model.batch_encode(docs)
|
||||
|
||||
data = [{
|
||||
@ -74,7 +76,9 @@ class KnowledgeBase:
|
||||
"vector": vectors[i],
|
||||
"text": docs[i],
|
||||
"hash": hashstr(docs[i], with_salt=True),
|
||||
**kwargs} for i in range(len(vectors))]
|
||||
**kwargs,
|
||||
**chunk_infos[i]
|
||||
} for i in range(len(vectors))]
|
||||
|
||||
res = self.client.insert(collection_name=collection_name, data=data)
|
||||
return res
|
||||
|
||||
@ -94,8 +94,8 @@
|
||||
@after-open-change="afterOpenChange"
|
||||
>
|
||||
<h2>共 {{ selectedFile?.lines.length }} 个片段</h2>
|
||||
<p v-for="line in selectedFile?.lines" :key="line.id">
|
||||
<strong>Chunk #{{ line.id }}</strong> {{ line.text }}
|
||||
<p v-for="line in selectedFile?.lines" :key="line.id" class="line-text">
|
||||
{{ line.text }}
|
||||
</p>
|
||||
</a-drawer>
|
||||
</div>
|
||||
@ -871,6 +871,15 @@ onUnmounted(() => {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
.custom-class .line-text {
|
||||
padding: 10px;
|
||||
border-radius: 4px;
|
||||
|
||||
&:hover {
|
||||
background-color: var(--main-light-4);
|
||||
}
|
||||
}
|
||||
</style>
|
||||
|
||||
<style lang="less">
|
||||
|
||||
Loading…
Reference in New Issue
Block a user