[全量] 初始化项目代码、配置、文档及Agent协同harness
This commit is contained in:
580
langchain-chat/server/knowledge_base/utils.py
Normal file
580
langchain-chat/server/knowledge_base/utils.py
Normal file
@@ -0,0 +1,580 @@
|
||||
import asyncio
|
||||
import os
|
||||
import re
|
||||
from configs import (
|
||||
KB_ROOT_PATH,
|
||||
CHUNK_SIZE,
|
||||
OVERLAP_SIZE,
|
||||
ZH_TITLE_ENHANCE,
|
||||
logger,
|
||||
log_verbose,
|
||||
text_splitter_dict,
|
||||
LLM_MODELS,
|
||||
TEXT_SPLITTER_NAME,
|
||||
TEXT_SPLITTER_MAP
|
||||
)
|
||||
import importlib
|
||||
from server.chat.policy_fun_iast import get_llm_model_response_async
|
||||
from server.knowledge_base import kb_service as tr
|
||||
from server.knowledge_base.TexkRank import TextRank
|
||||
from text_splitter import zh_title_enhance as func_zh_title_enhance
|
||||
import langchain.document_loaders
|
||||
from langchain.docstore.document import Document
|
||||
from langchain.text_splitter import TextSplitter
|
||||
from pathlib import Path
|
||||
from server.utils import run_in_thread_pool, get_model_worker_config
|
||||
import json
|
||||
from typing import List, Union,Dict, Tuple, Generator
|
||||
import chardet
|
||||
import time
|
||||
|
||||
def get_split_time(f):
|
||||
def inner(*arg,**kwarg):
|
||||
s_time = time.time()
|
||||
res = f(*arg,**kwarg)
|
||||
e_time = time.time()
|
||||
print('切片耗时：{}秒'.format(e_time - s_time))
|
||||
return res
|
||||
return inner
|
||||
|
||||
def validate_kb_name(knowledge_base_id: str) -> bool:
|
||||
# 检查是否包含预期外的字符或路径攻击关键字
|
||||
if "../" in knowledge_base_id:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def get_kb_path(knowledge_base_name: str):
|
||||
return os.path.join(KB_ROOT_PATH, knowledge_base_name)
|
||||
|
||||
|
||||
def get_doc_path(knowledge_base_name: str):
|
||||
return os.path.join(get_kb_path(knowledge_base_name), "content")
|
||||
|
||||
|
||||
def get_vs_path(knowledge_base_name: str, vector_name: str):
|
||||
return os.path.join(get_kb_path(knowledge_base_name), "vector_store", vector_name)
|
||||
|
||||
|
||||
def get_file_path(knowledge_base_name: str, doc_name: str):
|
||||
return os.path.join(get_doc_path(knowledge_base_name), doc_name)
|
||||
|
||||
|
||||
def list_kbs_from_folder():
|
||||
return [f for f in os.listdir(KB_ROOT_PATH)
|
||||
if os.path.isdir(os.path.join(KB_ROOT_PATH, f))]
|
||||
|
||||
|
||||
def list_files_from_folder(kb_name: str):
|
||||
doc_path = get_doc_path(kb_name)
|
||||
result = []
|
||||
|
||||
def is_skiped_path(path: str):
|
||||
tail = os.path.basename(path).lower()
|
||||
for x in ["temp", "tmp", ".", "~$"]:
|
||||
if tail.startswith(x):
|
||||
return True
|
||||
return False
|
||||
|
||||
def process_entry(entry):
|
||||
if is_skiped_path(entry.path):
|
||||
return
|
||||
|
||||
if entry.is_symlink():
|
||||
target_path = os.path.realpath(entry.path)
|
||||
with os.scandir(target_path) as target_it:
|
||||
for target_entry in target_it:
|
||||
process_entry(target_entry)
|
||||
elif entry.is_file():
|
||||
file_path = (Path(os.path.relpath(entry.path, doc_path)).as_posix()) # 路径统一为 posix 格式
|
||||
result.append(file_path)
|
||||
elif entry.is_dir():
|
||||
with os.scandir(entry.path) as it:
|
||||
for sub_entry in it:
|
||||
process_entry(sub_entry)
|
||||
|
||||
with os.scandir(doc_path) as it:
|
||||
for entry in it:
|
||||
process_entry(entry)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
LOADER_DICT = {"GCYHTMLLoader": ['.html', '.htm'],
|
||||
"GCYWordLoader2": ['.docx', '.doc'],
|
||||
# "GCYWordLoader": ['.docx'], # .doc 解析目前有点问题，暂时关掉
|
||||
"MHTMLLoader": ['.mhtml'],
|
||||
"TextLoader": ['.md', '.txt'],
|
||||
"JSONLoader": [".json"],
|
||||
"JSONLinesLoader": [".jsonl"],
|
||||
"RapidOCRCSVLoader": [".csv"],
|
||||
# "CSVLoader": [".csv"],
|
||||
# "FilteredCSVLoader": [".csv"], 如果使用自定义分割csv
|
||||
"PyMuPDFLoader": [".pdf"],
|
||||
#"RapidOCRDocLoader": ['.docx', '.doc'],
|
||||
"RapidOCRPPTLoader": ['.ppt', '.pptx', ],
|
||||
"RapidOCRLoader": ['.png', '.jpg', '.jpeg', '.bmp'],
|
||||
"UnstructuredFileLoader": ['.eml', '.msg', '.rst',
|
||||
'.rtf', '.xml',
|
||||
'.epub', '.odt','.tsv'],
|
||||
"UnstructuredEmailLoader": ['.eml', '.msg'],
|
||||
"UnstructuredEPubLoader": ['.epub'],
|
||||
"ExcelLoader": ['.xlsx', '.xls', '.xlsd'],
|
||||
"NotebookLoader": ['.ipynb'],
|
||||
"UnstructuredODTLoader": ['.odt'],
|
||||
"PythonLoader": ['.py'],
|
||||
"UnstructuredRSTLoader": ['.rst'],
|
||||
"UnstructuredRTFLoader": ['.rtf'],
|
||||
"SRTLoader": ['.srt'],
|
||||
"TomlLoader": ['.toml'],
|
||||
"UnstructuredTSVLoader": ['.tsv'],
|
||||
"UnstructuredXMLLoader": ['.xml'],
|
||||
"UnstructuredPowerPointLoader": ['.ppt', '.pptx'],
|
||||
"EverNoteLoader": ['.enex'],
|
||||
}
|
||||
SUPPORTED_EXTS = [ext for sublist in LOADER_DICT.values() for ext in sublist]
|
||||
|
||||
|
||||
# patch json.dumps to disable ensure_ascii
|
||||
def _new_json_dumps(obj, **kwargs):
|
||||
kwargs["ensure_ascii"] = False
|
||||
return _origin_json_dumps(obj, **kwargs)
|
||||
|
||||
if json.dumps is not _new_json_dumps:
|
||||
_origin_json_dumps = json.dumps
|
||||
json.dumps = _new_json_dumps
|
||||
|
||||
|
||||
class JSONLinesLoader(langchain.document_loaders.JSONLoader):
|
||||
'''
|
||||
行式 Json 加载器，要求文件扩展名为 .jsonl
|
||||
'''
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self._json_lines = True
|
||||
|
||||
|
||||
langchain.document_loaders.JSONLinesLoader = JSONLinesLoader
|
||||
|
||||
|
||||
def get_LoaderClass(file_extension):
|
||||
for LoaderClass, extensions in LOADER_DICT.items():
|
||||
if file_extension in extensions:
|
||||
return LoaderClass
|
||||
|
||||
def get_SplitterClass(file_extension):
|
||||
"""
|
||||
根据文件类型获取文本分块器类型
|
||||
"""
|
||||
# print('get Splitter Class', file_extension)
|
||||
for SplitterClass, extensions in TEXT_SPLITTER_MAP.items():
|
||||
if file_extension in extensions:
|
||||
return SplitterClass
|
||||
print(f'未找到文件类型"{file_extension}"对应的切分器，使用默认值')
|
||||
return TEXT_SPLITTER_NAME
|
||||
|
||||
def get_loader(loader_name: str, file_path: str, loader_kwargs: Dict = None):
|
||||
'''
|
||||
根据loader_name和文件路径或内容返回文档加载器。
|
||||
'''
|
||||
loader_kwargs = loader_kwargs or {}
|
||||
try:
|
||||
# print(loader_name)
|
||||
if loader_name in ["RapidOCRLoader", "FilteredCSVLoader",
|
||||
"GCYWordLoader","GCYWordLoader2", "GCYHTMLLoader",
|
||||
"RapidOCRPPTLoader", "RapidOCRCSVLoader","ExcelLoader"]:
|
||||
|
||||
document_loaders_module = importlib.import_module('document_loaders')
|
||||
else:
|
||||
document_loaders_module = importlib.import_module('langchain.document_loaders')
|
||||
DocumentLoader = getattr(document_loaders_module, loader_name)
|
||||
except Exception as e:
|
||||
msg = f"为文件{file_path}查找加载器{loader_name}时出错：{e}"
|
||||
logger.error(f'{e.__class__.__name__}: {msg}',
|
||||
exc_info=e if log_verbose else None)
|
||||
document_loaders_module = importlib.import_module('langchain.document_loaders')
|
||||
DocumentLoader = getattr(document_loaders_module, "UnstructuredFileLoader")
|
||||
|
||||
if loader_name in ["UnstructuredFileLoader", "TextLoader"]:
|
||||
loader_kwargs.setdefault("autodetect_encoding", True)
|
||||
elif loader_name == "CSVLoader":
|
||||
if not loader_kwargs.get("encoding"):
|
||||
# 如果未指定 encoding，自动识别文件编码类型，避免langchain loader 加载文件报编码错误
|
||||
with open(file_path, 'rb') as struct_file:
|
||||
encode_detect = chardet.detect(struct_file.read())
|
||||
if encode_detect is None:
|
||||
encode_detect = {"encoding": "utf-8"}
|
||||
loader_kwargs["encoding"] = encode_detect["encoding"]
|
||||
|
||||
elif loader_name == "JSONLoader":
|
||||
loader_kwargs.setdefault("jq_schema", ".")
|
||||
loader_kwargs.setdefault("text_content", False)
|
||||
elif loader_name == "JSONLinesLoader":
|
||||
loader_kwargs.setdefault("jq_schema", ".")
|
||||
loader_kwargs.setdefault("text_content", False)
|
||||
|
||||
loader = DocumentLoader(file_path, **loader_kwargs)
|
||||
return loader
|
||||
|
||||
|
||||
def make_text_splitter(
|
||||
splitter_name: str = TEXT_SPLITTER_NAME,
|
||||
chunk_size: int = CHUNK_SIZE,
|
||||
chunk_overlap: int = OVERLAP_SIZE,
|
||||
llm_model: str = LLM_MODELS[0],
|
||||
):
|
||||
"""
|
||||
根据参数获取特定的分词器
|
||||
"""
|
||||
# print('spliter name', splitter_name)
|
||||
splitter_name = splitter_name or "SpacyTextSplitter"
|
||||
try:
|
||||
# if splitter_name == "GCYMarkdownTextSplitter": # MarkdownHeaderTextSplitter特殊判定
|
||||
if splitter_name == "MarkdownTextSplitter": # MarkdownHeaderTextSplitter特殊判定
|
||||
text_splitter_module = importlib.import_module('text_splitter')
|
||||
TextSplitter = getattr(text_splitter_module, splitter_name)
|
||||
headers_to_split_on = text_splitter_dict[splitter_name]['headers_to_split_on']
|
||||
text_splitter = TextSplitter(
|
||||
headers_to_split_on=headers_to_split_on,
|
||||
strip_headers=False, # 不要将标题从分块文本中去掉
|
||||
promote_headers=True
|
||||
)
|
||||
else:
|
||||
|
||||
try: ## 优先使用用户自定义的text_splitter
|
||||
text_splitter_module = importlib.import_module('text_splitter')
|
||||
TextSplitter = getattr(text_splitter_module, splitter_name)
|
||||
except: ## 否则使用langchain的text_splitter
|
||||
text_splitter_module = importlib.import_module('langchain.text_splitter')
|
||||
TextSplitter = getattr(text_splitter_module, splitter_name)
|
||||
|
||||
if text_splitter_dict[splitter_name]["source"] == "tiktoken": ## 从tiktoken加载
|
||||
try:
|
||||
text_splitter = TextSplitter.from_tiktoken_encoder(
|
||||
encoding_name=text_splitter_dict[splitter_name]["tokenizer_name_or_path"],
|
||||
pipeline="zh_core_web_sm",
|
||||
chunk_size=chunk_size,
|
||||
chunk_overlap=chunk_overlap
|
||||
)
|
||||
except:
|
||||
text_splitter = TextSplitter.from_tiktoken_encoder(
|
||||
encoding_name=text_splitter_dict[splitter_name]["tokenizer_name_or_path"],
|
||||
chunk_size=chunk_size,
|
||||
chunk_overlap=chunk_overlap
|
||||
)
|
||||
elif text_splitter_dict[splitter_name]["source"] == "huggingface": ## 从huggingface加载
|
||||
if text_splitter_dict[splitter_name]["tokenizer_name_or_path"] == "":
|
||||
config = get_model_worker_config(llm_model)
|
||||
text_splitter_dict[splitter_name]["tokenizer_name_or_path"] = \
|
||||
config.get("model_path")
|
||||
|
||||
if text_splitter_dict[splitter_name]["tokenizer_name_or_path"] == "gpt2":
|
||||
from transformers import GPT2TokenizerFast
|
||||
from langchain.text_splitter import CharacterTextSplitter
|
||||
tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
|
||||
else: ## 字符长度加载
|
||||
from transformers import AutoTokenizer
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
text_splitter_dict[splitter_name]["tokenizer_name_or_path"],
|
||||
trust_remote_code=True)
|
||||
text_splitter = TextSplitter.from_huggingface_tokenizer(
|
||||
tokenizer=tokenizer,
|
||||
chunk_size=chunk_size,
|
||||
chunk_overlap=chunk_overlap
|
||||
)
|
||||
elif text_splitter_dict[splitter_name]["source"] == "no_tokenizer": # IAST 0429: 目前不需要使用分词器
|
||||
text_splitter = TextSplitter(
|
||||
chunk_size=chunk_size,
|
||||
chunk_overlap=chunk_overlap
|
||||
)
|
||||
else:
|
||||
try:
|
||||
text_splitter = TextSplitter(
|
||||
pipeline="zh_core_web_sm",
|
||||
chunk_size=chunk_size,
|
||||
chunk_overlap=chunk_overlap
|
||||
)
|
||||
except:
|
||||
text_splitter = TextSplitter(
|
||||
chunk_size=chunk_size,
|
||||
chunk_overlap=chunk_overlap
|
||||
)
|
||||
except Exception as e:
|
||||
print(e)
|
||||
text_splitter_module = importlib.import_module('langchain.text_splitter')
|
||||
TextSplitter = getattr(text_splitter_module, "RecursiveCharacterTextSplitter")
|
||||
text_splitter = TextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
|
||||
|
||||
# If you use SpacyTextSplitter you can use GPU to do split likes Issue #1287
|
||||
# text_splitter._tokenizer.max_length = 37016792
|
||||
# text_splitter._tokenizer.prefer_gpu()
|
||||
return text_splitter
|
||||
|
||||
|
||||
class KnowledgeFile:
|
||||
def __init__(
|
||||
self,
|
||||
filename: str,
|
||||
knowledge_base_name: str,
|
||||
loader_kwargs: Dict = {},
|
||||
):
|
||||
'''
|
||||
对应知识库目录中的文件，必须是磁盘上存在的才能进行向量化等操作。
|
||||
'''
|
||||
self.kb_name = knowledge_base_name
|
||||
self.filename = str(Path(filename).as_posix())
|
||||
self.ext = os.path.splitext(filename)[-1].lower()
|
||||
if self.ext not in SUPPORTED_EXTS:
|
||||
raise ValueError(f"暂未支持的文件格式 {self.filename}")
|
||||
self.loader_kwargs = loader_kwargs
|
||||
self.filepath = get_file_path(knowledge_base_name, filename)
|
||||
self.docs = None
|
||||
self.splited_docs = None
|
||||
self.document_loader_name = get_LoaderClass(self.ext)
|
||||
self.text_splitter_name = get_SplitterClass(self.ext)
|
||||
def get_full_text(self) -> Dict[str, str]:
|
||||
"""
|
||||
获取文件的全文内容，并返回文件名和全文内容的结构。
|
||||
"""
|
||||
try:
|
||||
docs = self.file2docs()
|
||||
full_text = "".join([doc.page_content for doc in docs])
|
||||
result = json.dumps( {
|
||||
"filename": self.filename,
|
||||
"full_text": full_text
|
||||
}, ensure_ascii=False)
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error(f"获取文件全文内容时出错：{e}", exc_info=e if log_verbose else None)
|
||||
return {
|
||||
"filename": self.filename,
|
||||
"full_text": "加载文件失败或文件内容为空"
|
||||
}
|
||||
async def get_llm_result(self) -> Dict[str, str]:
|
||||
"""
|
||||
根据文件的全文内容，异步调用模型生成文章摘要、关键词和章节速览。
|
||||
"""
|
||||
try:
|
||||
# full_text_data = self.get_full_text()
|
||||
# full_text = full_text_data.get("full_text", "")
|
||||
loop = asyncio.get_event_loop()
|
||||
full_text_data = await loop.run_in_executor(None, self.get_full_text)
|
||||
# full_text = full_text_data.get("full_text", "")
|
||||
try:
|
||||
# 将 JSON 字符串解析为字典
|
||||
full_text_dict = json.loads(full_text_data)
|
||||
full_text = full_text_dict.get("full_text", "")
|
||||
except json.JSONDecodeError:
|
||||
print("解析 JSON 数据时出错")
|
||||
full_text = ""
|
||||
if len(full_text) > 40000:
|
||||
# 判断英文占比
|
||||
# english_chars = re.findall(r'[a-zA-Z]', full_text)
|
||||
# english_ratio = len(english_chars) / len(full_text) if len(full_text) > 0 else 0
|
||||
# if english_ratio > 0.9 and len(full_text) > 50000:
|
||||
full_text = full_text[:40000]
|
||||
|
||||
# logger.info(f'=============文章长度{len(full_text)}')
|
||||
# full_text_80 = TextRank(full_text, 80)
|
||||
# logger.info(f'=============按80句压缩后文章长度{len(full_text_80)}')
|
||||
# if len(full_text_80) > 55000:
|
||||
# full_text_10 = TextRank(full_text_80, num_sentences=10)
|
||||
# logger.info(f'=============按10句压缩后文章长度{len(full_text_10)}')
|
||||
# full_text = full_text_10
|
||||
# else:
|
||||
# full_text = full_text_80
|
||||
else:
|
||||
pass
|
||||
# 异步调用模型
|
||||
from asyncio import gather
|
||||
llm_time = time.time()
|
||||
abstract_task = get_llm_model_response_async(
|
||||
strategy_name="gen_abstract",
|
||||
llm_model_name=LLM_MODELS[1],
|
||||
template_prompt_name="gen_abstract",
|
||||
prompt_param_dict={"context": full_text},
|
||||
temperature=0.7,
|
||||
max_tokens=4096
|
||||
)
|
||||
|
||||
keywords_task = get_llm_model_response_async(
|
||||
strategy_name="gen_keywords",
|
||||
llm_model_name=LLM_MODELS[1],
|
||||
template_prompt_name="gen_keywords",
|
||||
prompt_param_dict={"context": full_text},
|
||||
temperature=0.7,
|
||||
max_tokens=512
|
||||
)
|
||||
|
||||
paragraph_task = get_llm_model_response_async(
|
||||
strategy_name="gen_paragraph",
|
||||
llm_model_name=LLM_MODELS[0],
|
||||
template_prompt_name="gen_paragraph",
|
||||
prompt_param_dict={"context": full_text},
|
||||
temperature=0.7,
|
||||
max_tokens=8192
|
||||
)
|
||||
|
||||
# 并行执行任务
|
||||
article_abstract, article_keywords, article_paragraph = await gather(
|
||||
abstract_task, keywords_task, paragraph_task
|
||||
)
|
||||
|
||||
logger.info(f'生成导读用时：{time.time() - llm_time}')
|
||||
return {
|
||||
"filename": self.filename,
|
||||
"full_text": full_text,
|
||||
"article_abstract": article_abstract,
|
||||
"article_keywords": article_keywords,
|
||||
"article_paragraph": article_paragraph
|
||||
}
|
||||
except Exception as e:
|
||||
logger.error(f"生成LLM结果时出错：{e}", exc_info=e if log_verbose else None)
|
||||
return {
|
||||
"filename": self.filename,
|
||||
"article_abstract": "生成摘要失败",
|
||||
"article_keywords": "生成关键词失败",
|
||||
"article_paragraph": "生成章节速览失败"
|
||||
}
|
||||
def file2docs(self, refresh: bool = False):
|
||||
if self.docs is None or refresh:
|
||||
try:
|
||||
logger.info(f"{self.document_loader_name} used for {self.filepath}")
|
||||
loader = get_loader(loader_name=self.document_loader_name,
|
||||
file_path=self.filepath,
|
||||
loader_kwargs=self.loader_kwargs)
|
||||
self.docs = loader.load()
|
||||
except Exception as e:
|
||||
if self.document_loader_name == 'GCYWordLoader':
|
||||
loader = get_loader(loader_name='GCYWordLoader2',
|
||||
file_path=self.filepath,
|
||||
loader_kwargs=self.loader_kwargs)
|
||||
else:
|
||||
logger.error(f"加载文件 {self.filepath} 时出错：{e}", exc_info=e if log_verbose else None)
|
||||
self.docs = loader.load()
|
||||
return self.docs
|
||||
|
||||
@get_split_time
|
||||
def docs2texts(
|
||||
self,
|
||||
docs: List[Document] = None,
|
||||
zh_title_enhance: bool = ZH_TITLE_ENHANCE,
|
||||
refresh: bool = False,
|
||||
chunk_size: int = CHUNK_SIZE,
|
||||
chunk_overlap: int = OVERLAP_SIZE,
|
||||
text_splitter: TextSplitter = None,
|
||||
):
|
||||
docs = docs or self.file2docs(refresh=refresh)
|
||||
# debug 0429
|
||||
# print('docs2texts',docs )
|
||||
if not docs:
|
||||
return []
|
||||
|
||||
if text_splitter is None:
|
||||
self.text_splitter_name = get_SplitterClass(self.ext)
|
||||
text_splitter = make_text_splitter(splitter_name=self.text_splitter_name, chunk_size=chunk_size,
|
||||
chunk_overlap=chunk_overlap)
|
||||
# if self.text_splitter_name == "GCYMarkdownTextSplitter":
|
||||
if self.text_splitter_name == "MarkdownTextSplitter":
|
||||
doc_source = (docs[0].metadata)["source"]
|
||||
docs = text_splitter.split_markdown_text(docs[0].page_content, doc_source)
|
||||
else:
|
||||
docs = text_splitter.split_documents(docs)
|
||||
|
||||
if not docs:
|
||||
return []
|
||||
|
||||
# 检查切分好的文档是否有'h1'标题字段，如果没有，就加上。为之后入库其它有h1的文件做准备
|
||||
if 'h1' not in docs[0].metadata:
|
||||
for doc in docs:
|
||||
doc.metadata['h1'] = ''
|
||||
|
||||
print(f"文档切分示例：{docs[0]}")
|
||||
if zh_title_enhance:
|
||||
docs = func_zh_title_enhance(docs)
|
||||
self.splited_docs = docs
|
||||
return self.splited_docs
|
||||
|
||||
def file2text(
|
||||
self,
|
||||
zh_title_enhance: bool = ZH_TITLE_ENHANCE,
|
||||
refresh: bool = False,
|
||||
chunk_size: int = CHUNK_SIZE,
|
||||
chunk_overlap: int = OVERLAP_SIZE,
|
||||
text_splitter: TextSplitter = None,
|
||||
):
|
||||
if self.splited_docs is None or refresh:
|
||||
docs = self.file2docs()
|
||||
self.splited_docs = self.docs2texts(docs=docs,
|
||||
zh_title_enhance=zh_title_enhance,
|
||||
refresh=refresh,
|
||||
chunk_size=chunk_size,
|
||||
chunk_overlap=chunk_overlap,
|
||||
text_splitter=text_splitter)
|
||||
return self.splited_docs
|
||||
|
||||
def file_exist(self):
|
||||
return os.path.isfile(self.filepath)
|
||||
|
||||
def get_mtime(self):
|
||||
return os.path.getmtime(self.filepath)
|
||||
|
||||
def get_size(self):
|
||||
return os.path.getsize(self.filepath)
|
||||
|
||||
|
||||
def files2docs_in_thread(
|
||||
files: List[Union[KnowledgeFile, Tuple[str, str], Dict]],
|
||||
chunk_size: int = CHUNK_SIZE,
|
||||
chunk_overlap: int = OVERLAP_SIZE,
|
||||
zh_title_enhance: bool = ZH_TITLE_ENHANCE,
|
||||
) -> Generator:
|
||||
'''
|
||||
利用多线程批量将磁盘文件转化成langchain Document.
|
||||
如果传入参数是Tuple，形式为(filename, kb_name)
|
||||
生成器返回值为 status, (kb_name, file_name, docs | error)
|
||||
'''
|
||||
|
||||
def file2docs(*, file: KnowledgeFile, **kwargs) -> Tuple[bool, Tuple[str, str, List[Document]]]:
|
||||
try:
|
||||
return True, (file.kb_name, file.filename, file.file2text(**kwargs))
|
||||
except Exception as e:
|
||||
msg = f"从文件 {file.kb_name}/{file.filename} 加载文档时出错：{e}"
|
||||
logger.error(f'{e.__class__.__name__}: {msg}',
|
||||
exc_info=e if log_verbose else None)
|
||||
return False, (file.kb_name, file.filename, msg)
|
||||
|
||||
kwargs_list = []
|
||||
for i, file in enumerate(files):
|
||||
kwargs = {}
|
||||
try:
|
||||
if isinstance(file, tuple) and len(file) >= 2:
|
||||
filename = file[0]
|
||||
kb_name = file[1]
|
||||
file = KnowledgeFile(filename=filename, knowledge_base_name=kb_name)
|
||||
elif isinstance(file, dict):
|
||||
filename = file.pop("filename")
|
||||
kb_name = file.pop("kb_name")
|
||||
kwargs.update(file)
|
||||
file = KnowledgeFile(filename=filename, knowledge_base_name=kb_name)
|
||||
kwargs["file"] = file
|
||||
kwargs["chunk_size"] = chunk_size
|
||||
kwargs["chunk_overlap"] = chunk_overlap
|
||||
kwargs["zh_title_enhance"] = zh_title_enhance
|
||||
kwargs_list.append(kwargs)
|
||||
except Exception as e:
|
||||
yield False, (kb_name, filename, str(e))
|
||||
|
||||
for result in run_in_thread_pool(func=file2docs, params=kwargs_list):
|
||||
yield result
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from pprint import pprint
|
||||
|
||||
kb_file = KnowledgeFile(
|
||||
filename="/home/congyin/Code/Project_Langchain_0814/Langchain-Chatchat/knowledge_base/csv1/content/gm.csv",
|
||||
knowledge_base_name="samples")
|
||||
# kb_file.text_splitter_name = "RecursiveCharacterTextSplitter"
|
||||
docs = kb_file.file2docs()
|
||||
# pprint(docs[-1])
|
||||
Reference in New Issue
Block a user