Files
gangyan/langchain-chat/server/knowledge_base/utils.py

581 lines
24 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

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[0],
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[0],
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])