Spaces:
Runtime error
Runtime error
# Copyright (c) OpenMMLab. All rights reserved. | |
"""extract feature and search with user query.""" | |
from __future__ import annotations | |
import argparse | |
import json | |
import os | |
import re | |
import shutil | |
# from multiprocessing import Pool | |
from typing import Any, List, Optional | |
import pytoml | |
from BCEmbedding.tools.langchain import BCERerank | |
from langchain_community.embeddings import HuggingFaceEmbeddings | |
from langchain.text_splitter import (MarkdownHeaderTextSplitter, | |
MarkdownTextSplitter, | |
RecursiveCharacterTextSplitter) | |
from langchain_community.vectorstores.faiss import FAISS as Vectorstore | |
from langchain_core.documents import Document | |
from loguru import logger | |
from torch.cuda import empty_cache | |
from .file_operation import FileName, FileOperation | |
from .retriever import CacheRetriever, Retriever | |
from .cluster import Clusterer | |
def save_image(image, path, name): | |
""" 保存图片到指定路径 """ | |
if not os.path.exists(path): | |
os.makedirs(path) | |
image_path = os.path.join(path, name) | |
image.save(image_path) | |
return image_path | |
def save_all_image(image_folder,tables): | |
# 假设 res 中包含了图片对象和其他数据 | |
for index, data in enumerate(tables): | |
image, text = data # 假设 data 结构是这样的 | |
image_path = save_image(image, image_folder, f'image_{index}.png') | |
# relative_path = os.path.relpath(image_path, preprocessdir) | |
tables[index] = (image_path, text) # 更新 res 中的图片对象为图片路径 | |
return tables | |
def create_html_file(tables, html_file_path): | |
html_content = '<html><body>\n' | |
for index, data in enumerate(tables): | |
image_path, text = data | |
# 创建图片链接和文本 | |
html_content += f'<img src="{image_path}" alt="Image">\n{text}\n' | |
html_content += '</body></html>' | |
# 写入 HTML 文件 | |
with open(html_file_path, 'w') as file: | |
file.write(html_content) | |
def read_and_save(file: FileName,file_opr: FileOperation): | |
if os.path.exists(file.copypath) and os.path.exists(file.jsonpath) and os.path.exists(file.htmlpath) and os.path.exists(file.imagefolder): | |
# already exists, return | |
logger.info('already exist, skip load') | |
return | |
# file_opr = FileOperation() | |
logger.info('reading {}, would save to {}'.format(file.origin, | |
file.copypath)) | |
content, tbls, error = file_opr.read(file.origin) | |
if error is not None: | |
logger.error('{} load error: {}'.format(file.origin, str(error))) | |
return | |
if content is None or len(content) < 1: | |
logger.warning('{} empty, skip save'.format(file.origin)) | |
return | |
with open(file.copypath, 'w') as f: | |
f.write(content) | |
tables = save_all_image(file.imagefolder,tbls) | |
with open(file.jsonpath, 'w') as f: | |
json.dump(tables, f, indent=4, ensure_ascii=False) | |
create_html_file(tables, file.htmlpath) | |
def _split_text_with_regex_from_end(text: str, separator: str, | |
keep_separator: bool) -> List[str]: | |
# Now that we have the separator, split the text | |
if separator: | |
if keep_separator: | |
# The parentheses in the pattern keep the delimiters in the result. | |
_splits = re.split(f'({separator})', text) | |
splits = [''.join(i) for i in zip(_splits[0::2], _splits[1::2])] | |
if len(_splits) % 2 == 1: | |
splits += _splits[-1:] | |
# splits = [_splits[0]] + splits | |
else: | |
splits = re.split(separator, text) | |
else: | |
splits = list(text) | |
return [s for s in splits if s != ''] | |
# copy from https://github.com/chatchat-space/Langchain-Chatchat/blob/master/text_splitter/chinese_recursive_text_splitter.py | |
class ChineseRecursiveTextSplitter(RecursiveCharacterTextSplitter): | |
def __init__( | |
self, | |
separators: Optional[List[str]] = None, | |
keep_separator: bool = True, | |
is_separator_regex: bool = True, | |
**kwargs: Any, | |
) -> None: | |
"""Create a new TextSplitter.""" | |
super().__init__(keep_separator=keep_separator, **kwargs) | |
self._separators = separators or [ | |
'\n\n', '\n', '。|!|?', '\.\s|\!\s|\?\s', ';|;\s', ',|,\s' | |
] | |
self._is_separator_regex = is_separator_regex | |
def _split_text(self, text: str, separators: List[str]) -> List[str]: | |
"""Split incoming text and return chunks.""" | |
final_chunks = [] | |
# Get appropriate separator to use | |
separator = separators[-1] | |
new_separators = [] | |
for i, _s in enumerate(separators): | |
_separator = _s if self._is_separator_regex else re.escape(_s) | |
if _s == '': | |
separator = _s | |
break | |
if re.search(_separator, text): | |
separator = _s | |
new_separators = separators[i + 1:] | |
break | |
_separator = separator if self._is_separator_regex else re.escape( | |
separator) | |
splits = _split_text_with_regex_from_end(text, _separator, | |
self._keep_separator) | |
# Now go merging things, recursively splitting longer texts. | |
_good_splits = [] | |
_separator = '' if self._keep_separator else separator | |
for s in splits: | |
if self._length_function(s) < self._chunk_size: | |
_good_splits.append(s) | |
else: | |
if _good_splits: | |
merged_text = self._merge_splits(_good_splits, _separator) | |
final_chunks.extend(merged_text) | |
_good_splits = [] | |
if not new_separators: | |
final_chunks.append(s) | |
else: | |
other_info = self._split_text(s, new_separators) | |
final_chunks.extend(other_info) | |
if _good_splits: | |
merged_text = self._merge_splits(_good_splits, _separator) | |
final_chunks.extend(merged_text) | |
return [ | |
re.sub(r'\n{2,}', '\n', chunk.strip()) for chunk in final_chunks | |
if chunk.strip() != '' | |
] | |
class FeatureStore: | |
"""Tokenize and extract features from the project's documents, for use in | |
the reject pipeline and response pipeline.""" | |
def __init__(self, | |
embeddings: HuggingFaceEmbeddings, | |
reranker: BCERerank, | |
chunk_size: int, | |
n_clusters: int|list[int], | |
config_path: str = 'config.ini', | |
language: str = 'zh') -> None: | |
"""Init with model device type and config.""" | |
self.config_path = config_path | |
self.reject_throttle = -1 | |
self.language = language | |
with open(config_path, encoding='utf8') as f: | |
config = pytoml.load(f)['feature_store'] | |
self.reject_throttle = config['reject_throttle'] | |
self.chunk_size = chunk_size | |
if isinstance(n_clusters, int): | |
self.n_clusters = [n_clusters] | |
elif isinstance(n_clusters, list): | |
self.n_clusters = n_clusters | |
logger.warning( | |
'!!! If your feature generated by `text2vec-large-chinese` before 20240208, please rerun `python3 -m huixiangdou.service.feature_store`' # noqa E501 | |
) | |
logger.debug('loading text2vec model..') | |
self.embeddings = embeddings | |
self.reranker = reranker | |
self.compression_retriever = None | |
self.rejecter = None | |
self.retriever = None | |
self.md_splitter = MarkdownTextSplitter(chunk_size=self.chunk_size, | |
chunk_overlap=32) | |
if language == 'zh': | |
self.text_splitter = ChineseRecursiveTextSplitter( | |
keep_separator=True, | |
is_separator_regex=True, | |
chunk_size=self.chunk_size, | |
chunk_overlap=32) | |
else: | |
self.text_splitter = RecursiveCharacterTextSplitter( | |
chunk_size=self.chunk_size, chunk_overlap=32) | |
self.head_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=[ | |
('#', 'Header 1'), | |
('##', 'Header 2'), | |
('###', 'Header 3'), | |
]) | |
def split_md(self, text: str, source: None): | |
"""Split the markdown document in a nested way, first extracting the | |
header. | |
If the extraction result exceeds 1024, split it again according to | |
length. | |
""" | |
docs = self.head_splitter.split_text(text) | |
final = [] | |
for doc in docs: | |
header = '' | |
if len(doc.metadata) > 0: | |
if 'Header 1' in doc.metadata: | |
header += doc.metadata['Header 1'] | |
if 'Header 2' in doc.metadata: | |
header += ' ' | |
header += doc.metadata['Header 2'] | |
if 'Header 3' in doc.metadata: | |
header += ' ' | |
header += doc.metadata['Header 3'] | |
if len(doc.page_content) >= 1024: | |
subdocs = self.md_splitter.create_documents([doc.page_content]) | |
for subdoc in subdocs: | |
if len(subdoc.page_content) >= 10: | |
final.append('{} {}'.format( | |
header, subdoc.page_content.lower())) | |
elif len(doc.page_content) >= 10: | |
final.append('{} {}'.format( | |
header, doc.page_content.lower())) # noqa E501 | |
for item in final: | |
if len(item) >= 1024: | |
logger.debug('source {} split length {}'.format( | |
source, len(item))) | |
return final | |
def clean_md(self, text: str): | |
"""Remove parts of the markdown document that do not contain the key | |
question words, such as code blocks, URL links, etc.""" | |
# remove ref | |
pattern_ref = r'\[(.*?)\]\(.*?\)' | |
new_text = re.sub(pattern_ref, r'\1', text) | |
# remove code block | |
pattern_code = r'```.*?```' | |
new_text = re.sub(pattern_code, '', new_text, flags=re.DOTALL) | |
# remove underline | |
new_text = re.sub('_{5,}', '', new_text) | |
# remove table | |
# new_text = re.sub('\|.*?\|\n\| *\:.*\: *\|.*\n(\|.*\|.*\n)*', '', new_text, flags=re.DOTALL) # noqa E501 | |
# use lower | |
new_text = new_text.lower() | |
return new_text | |
def get_md_documents(self, file: FileName): | |
documents = [] | |
length = 0 | |
text = '' | |
with open(file.copypath, encoding='utf8') as f: | |
text = f.read() | |
text = file.prefix + '\n' + self.clean_md(text) | |
if len(text) <= 1: | |
return [], length | |
chunks = self.split_md(text=text, | |
source=os.path.abspath(file.copypath)) | |
for chunk in chunks: | |
new_doc = Document(page_content=chunk, | |
metadata={ | |
'source': file.basename, | |
'read': file.copypath | |
}) | |
length += len(chunk) | |
documents.append(new_doc) | |
return documents, length | |
def get_text_documents(self, text: str, file: FileName): | |
if len(text) <= 1: | |
return [] | |
chunks = self.text_splitter.create_documents([text]) | |
documents = [] | |
for chunk in chunks: | |
# `source` is for return references | |
# `read` is for LLM response | |
chunk.metadata = {'source': file.basename, 'read': file.copypath} | |
documents.append(chunk) | |
return documents | |
def ingress_response(self, files: list, work_dir: str, file_opr: FileOperation): | |
"""Extract the features required for the response pipeline based on the | |
document.""" | |
chunk_dir = os.path.join(work_dir,f"chunksize_{self.chunk_size}") | |
feature_dir = os.path.join(chunk_dir, 'db_response') | |
if not os.path.exists(feature_dir): | |
os.makedirs(feature_dir) | |
# logger.info('glob {} in dir {}'.format(files, file_dir)) | |
# file_opr = FileOperation() | |
documents = [] | |
for i, file in enumerate(files): | |
logger.debug('{}/{}.. {}'.format(i + 1, len(files), file.basename)) | |
if not file.state: | |
continue | |
if file._type == 'md': | |
md_documents, md_length = self.get_md_documents(file) | |
documents += md_documents | |
logger.info('{} content length {}'.format( | |
file._type, md_length)) | |
file.reason = str(md_length) | |
else: | |
# now read pdf/word/excel/ppt text | |
text, _,error = file_opr.read(file.copypath) | |
if error is not None: | |
file.state = False | |
file.reason = str(error) | |
continue | |
file.reason = str(len(text)) | |
logger.info('{} content length {}'.format(file._type, len(text))) | |
text = file.prefix + text | |
documents += self.get_text_documents(text, file) | |
if len(documents) < 1: | |
return | |
texts = [doc.page_content for doc in documents] | |
logger.debug('calculating embeddings..') | |
text_embeddings = self.embeddings.embed_documents(texts) | |
logger.debug('clustering response data..') | |
clusterer = Clusterer(texts, text_embeddings, n_clusters= self.n_clusters) | |
clusterer.generate_cluster(chunk_dir) | |
text_embedding_pairs = zip(texts, text_embeddings) | |
metadatas = [doc.metadata for doc in documents] | |
logger.debug('making vectorstore..') | |
vs = Vectorstore.from_embeddings(text_embedding_pairs, self.embeddings, metadatas) | |
vs.save_local(feature_dir) | |
# def ingress_reject(self, files: list, work_dir: str): | |
# """Extract the features required for the reject pipeline based on | |
# documents.""" | |
# feature_dir = os.path.join(work_dir, 'db_reject') | |
# if not os.path.exists(feature_dir): | |
# os.makedirs(feature_dir) | |
# documents = [] | |
# file_opr = FileOperation() | |
# logger.debug('ingress reject..') | |
# for i, file in enumerate(files): | |
# if not file.state: | |
# continue | |
# if file._type == 'md': | |
# # reject base not clean md | |
# text = file.basename + '\n' | |
# with open(file.copypath, encoding='utf8') as f: | |
# text += f.read() | |
# if len(text) <= 1: | |
# continue | |
# chunks = self.split_md(text=text, | |
# source=os.path.abspath(file.copypath)) | |
# for chunk in chunks: | |
# new_doc = Document(page_content=chunk, | |
# metadata={ | |
# 'source': file.basename, | |
# 'read': file.copypath | |
# }) | |
# documents.append(new_doc) | |
# else: | |
# text, error = file_opr.read(file.copypath) | |
# if error is not None: | |
# continue | |
# text = file.basename + text | |
# documents += self.get_text_documents(text, file) | |
# if len(documents) < 1: | |
# return | |
# vs = Vectorstore.from_documents(documents, self.embeddings) | |
# vs.save_local(feature_dir) | |
def preprocess(self, files: list, work_dir: str,file_opr: FileOperation): | |
"""Preprocesses files in a given directory. Copies each file to | |
'preprocess' with new name formed by joining all subdirectories with | |
'_'. | |
Args: | |
files (list): original file list. | |
work_dir (str): Working directory where preprocessed files will be stored. # noqa E501 | |
Returns: | |
str: Path to the directory where preprocessed markdown files are saved. | |
Raises: | |
Exception: Raise an exception if no markdown files are found in the provided repository directory. # noqa E501 | |
""" | |
preproc_dir = os.path.join(work_dir, 'preprocess') | |
if not os.path.exists(preproc_dir): | |
os.makedirs(preproc_dir) | |
# pool = Pool(processes=16) | |
# file_opr = FileOperation() | |
logger.debug('waiting for preprocess read finish..') | |
for idx, file in enumerate(files): | |
if not os.path.exists(file.origin): | |
file.state = False | |
file.reason = 'skip not exist' | |
continue | |
if file._type == 'image': | |
file.state = False | |
file.reason = 'skip image' | |
elif file._type in ['pdf']: | |
# read pdf/word/excel file and save to text format | |
md5 = file_opr.md5(file.origin) | |
file.copypath = os.path.join(preproc_dir,'{}.text'.format(md5)) | |
file.jsonpath = os.path.join(preproc_dir,'{}.json'.format(md5)) | |
file.htmlpath = os.path.join(preproc_dir,'{}.html'.format(md5)) | |
file.imagefolder = os.path.join(preproc_dir,'{}_images'.format(md5)) | |
# pool.apply_async(read_and_save, (file, )) | |
read_and_save(file,file_opr) | |
elif file._type in ['md', 'text']: | |
# rename text files to new dir | |
md5 = file_opr.md5(file.origin) | |
file.copypath = os.path.join( | |
preproc_dir, | |
file.origin.replace('/', '_')[-84:]) | |
try: | |
shutil.copy(file.origin, file.copypath) | |
file.state = True | |
file.reason = 'preprocessed' | |
except Exception as e: | |
file.state = False | |
file.reason = str(e) | |
else: # 'word', 'excel', 'ppt', 'html' TODO | |
file.state = False | |
file.reason = 'skip unknown format' | |
# pool.close() | |
# pool.join() | |
# check process result | |
for file in files: | |
if file._type in ['pdf', 'word', 'excel']: | |
if os.path.exists(file.copypath): | |
file.state = True | |
file.reason = 'preprocessed' | |
else: | |
file.state = False | |
file.reason = 'read error' | |
def saveconfig(self,workdir:str): | |
# shutil.copy(self.config_path, os.path.join(workdir, self.config_path)) | |
with open(self.config_path, encoding='utf8') as f: | |
config = pytoml.load(f) | |
config['feature_store']['n_clusters'] = self.n_clusters | |
config['feature_store']['chunk_size'] = self.chunk_size | |
chunk_dir = os.path.join(workdir,f"chunksize_{self.chunk_size}") | |
with open(os.path.join(workdir, self.config_path), 'w') as f: | |
pytoml.dump(config, f) | |
f.close() | |
with open(os.path.join(chunk_dir, self.config_path), 'w') as f: | |
pytoml.dump(config, f) | |
f.close() | |
with open(self.config_path, 'w') as f1: | |
pytoml.dump(config, f1) | |
f1.close() | |
def initialize(self, files: list, work_dir: str,file_opr: FileOperation): | |
"""Initializes response and reject feature store. | |
Only needs to be called once. Also calculates the optimal threshold | |
based on provided good and bad question examples, and saves it in the | |
configuration file. | |
""" | |
logger.info( | |
'initialize response and reject feature store, you only need call this once.' # noqa E501 | |
) | |
self.preprocess(files=files, work_dir=work_dir, file_opr=file_opr) | |
self.ingress_response(files=files, work_dir=work_dir, file_opr=file_opr) | |
# self.ingress_reject(files=files, work_dir=work_dir) | |
self.saveconfig(work_dir) | |
def parse_args(): | |
"""Parse command-line arguments.""" | |
parser = argparse.ArgumentParser( | |
description='Feature store for processing directories.') | |
parser.add_argument('--work_dir', | |
type=str, | |
default='workdir', | |
help='Working directory.') | |
parser.add_argument( | |
'--repo_dir', | |
type=str, | |
default='repodir', | |
help='Root directory where the repositories are located.') | |
parser.add_argument( | |
'--config_path', | |
default='config.ini', | |
help='Feature store configuration path. Default value is config.ini') | |
parser.add_argument( | |
'--good_questions', | |
default='resource/good_questions.json', | |
help= # noqa E251 | |
'Positive examples in the dataset. Default value is resource/good_questions.json' # noqa E501 | |
) | |
parser.add_argument( | |
'--bad_questions', | |
default='resource/bad_questions.json', | |
help= # noqa E251 | |
'Negative examples json path. Default value is resource/bad_questions.json' # noqa E501 | |
) | |
parser.add_argument( | |
'--sample', help='Input an json file, save reject and search output.') | |
args = parser.parse_args() | |
return args | |
def test_reject(retriever: Retriever, sample: str = None, work_dir: str = 'workdir'): | |
"""Simple test reject pipeline.""" | |
if sample is None: | |
real_questions = [ | |
'SAM 10个T 的训练集,怎么比比较公平呢~?速度上还有缺陷吧?', | |
'想问下,如果只是推理的话,amp的fp16是不会省显存么,我看parameter仍然是float32,开和不开推理的显存占用都是一样的。能不能直接用把数据和model都 .half() 代替呢,相比之下amp好在哪里', # noqa E501 | |
'mmdeploy支持ncnn vulkan部署么,我只找到了ncnn cpu 版本', | |
'大佬们,如果我想在高空检测安全帽,我应该用 mmdetection 还是 mmrotate', | |
'请问 ncnn 全称是什么', | |
'有啥中文的 text to speech 模型吗?', | |
'今天中午吃什么?', | |
'huixiangdou 是什么?', | |
'mmpose 如何安装?', | |
'使用科研仪器需要注意什么?' | |
] | |
else: | |
with open(sample) as f: | |
real_questions = json.load(f) | |
for example in real_questions: | |
reject, _ = retriever.is_reject(example) | |
if reject: | |
logger.error(f'reject query: {example}') | |
else: | |
logger.warning(f'process query: {example}') | |
if sample is not None: | |
if reject: | |
with open(f'{work_dir}/negative.txt', 'a+') as f: | |
f.write(example) | |
f.write('\n') | |
else: | |
with open(f'{work_dir}/positive.txt', 'a+') as f: | |
f.write(example) | |
f.write('\n') | |
empty_cache() | |
def test_query(retriever: Retriever, sample: str = None, work_dir: str = 'workdir'): | |
"""Simple test response pipeline.""" | |
if sample is not None: | |
with open(sample) as f: | |
real_questions = json.load(f) | |
logger.add('logs/feature_store_query.log', rotation='4MB') | |
else: | |
real_questions = ['mmpose installation', 'how to use std::vector ?'] | |
for example in real_questions: | |
example = example[0:400] | |
print(retriever.query(example)) | |
empty_cache() | |
empty_cache() | |
if __name__ == '__main__': | |
args = parse_args() | |
cache = CacheRetriever(config_path=args.config_path) | |
fs_init = FeatureStore(embeddings=cache.embeddings, | |
reranker=cache.reranker, | |
config_path=args.config_path) | |
with open(args.config_path, encoding='utf8') as f: | |
config = pytoml.load(f) | |
# walk all files in repo dir | |
file_opr = FileOperation() | |
files = file_opr.scan_dir(repo_dir=config['feature_store']['repo_dir']) | |
fs_init.initialize(files=files, work_dir=config['feature_store']['work_dir']) | |
file_opr.summarize(files) | |
del fs_init | |
# update reject throttle | |
retriever = cache.get(config_path=args.config_path, work_dir=config['feature_store']['work_dir']) | |
with open(os.path.join('resource', 'good_questions.json')) as f: | |
good_questions = json.load(f) | |
with open(os.path.join('resource', 'bad_questions.json')) as f: | |
bad_questions = json.load(f) | |
retriever.update_throttle(config_path=args.config_path, | |
good_questions=good_questions, | |
bad_questions=bad_questions) | |
cache.pop('default') | |
# test | |
retriever = cache.get(config_path=args.config_path, work_dir=config['feature_store']['work_dir']) | |
test_reject(retriever, args.sample, config['feature_store']['work_dir']) | |
test_query(retriever, args.sample, config['feature_store']['work_dir']) | |