from Reader.docx_reader import DocxReader from Reader.pdf_reader import PdfReader import threading import time import os from models import deepseek_online class Reader: def __init__(self, workers:int=2, task_max_count:int=100): # 任务列表 self.task_list = [] # 任务列表 self.finish_list = [] # 完成列表 # 互斥锁,注意不可重入 self.__lock = threading.Lock() # 条件变量 self.__condition = threading.Condition(self.__lock) # 任务列表最大元素 self.__list_max_limit = task_max_count # 开启线程 self.__threads = [] # 线程数量 self.workers = workers # 运行状态 self.__is_running = False # 守卫线程 self.__guard_thread = None self.__guard_period = 6 # 守护周期，单位：秒 # 预处理文档保存路径 self.tem_save_path = './tem_files' # 处理后的输出文件路径 self.chat_output_path = './chat_output' def start(self): """ 启动 Returns: """ # 先给他把路径创建上，免得用到的地方报错 # 创建临时路径，存放预处理文档 if not os.path.exists(self.tem_save_path): os.mkdir(self.tem_save_path) # 创建LLM输出路径 if not os.path.exists(self.chat_output_path): os.mkdir(self.chat_output_path) # 设置状态 with self.__condition: # 开启状态 self.__is_running = True for i in range(self.workers): print(f'start: 准备开启consumer线程：{i}') self.__threads.append(threading.Thread(target=self.consumer)) self.__threads[i].start() print("start: 准备开启守卫线程") self.__start_guard() def stop(self): """ 关闭 Returns: """ # 设置状态 with self.__condition: # 关闭状态 self.__is_running = False # 唤醒所有线程 self.__condition.notify_all() print('stop: 关闭时唤醒全部线程') # 释放锁 for i, t in enumerate(self.__threads): if t.is_alive(): t.join() print(f'stop: 等待守护线程结束, 预计时间{self.__guard_period}s') if self.__guard_thread.is_alive(): self.__guard_thread.join() print(f'stop: 所有线程均已退出') def __guard(self): """ 守护函数 Returns: """ prefixed = f'guard({threading.get_ident()}):' print(prefixed + "启动") while self.__is_running: running_thread = 0 for t in self.__threads: if t.is_alive(): running_thread += 1 else: print(prefixed+f"{t.ident}未存活") if running_thread != self.workers: print(prefixed + "有线程挂了！应该再启动一个") # 启动挂掉的线程 pass # 暂不实现 else: task_num = len(self.task_list) finish_num = len(os.listdir(self.chat_output_path)) # 通过读取已处理文件的个数来确定 print(prefixed + f"消费者线程运行数量: {running_thread}, 已处理任务数量:{finish_num}, 剩余任务数量: {task_num}") time.sleep(self.__guard_period) def __start_guard(self): self.__guard_thread = threading.Thread(target=self.__guard) self.__guard_thread.start() def add_task(self, task: str)->bool: """ 向任务列表添加任务 Args: task: 待处理的文件路径 Returns: 成功返回True，失败返回False """ flag = False with self.__condition: # 判断列表是否超限 if len(self.task_list) < self.__list_max_limit: # 未超限添加新任务 self.task_list.append(task) flag = True # 唤醒正在等待的线程1个 self.__condition.notify() print('add_task: 添加1个任务，唤醒1个线程') return flag def get_task_list_len(self)-> int: """ 获取任务列表长度 Returns: 整数 """ with self.__lock: task_list_len = len(self.task_list) return task_list_len def get_finish_list_len(self)->int: """ 获取完成列表长度 Returns: """ with self.__lock: finish_list_len = len(self.finish_list) return finish_list_len def consumer(self): """ 任务处理函数 Returns: """ thread_id = threading.get_ident() # 线程id prefixed = f'consumer({thread_id}):' print(prefixed + '开启消费者线程') while self.__is_running: with self.__condition: # 获取锁 while len(self.task_list) == 0 and self.__is_running: print(prefixed + '消费者进入等待') self.__condition.wait() # 进入等待 print(prefixed + '消费者被唤醒') if not self.__is_running: break # 被唤醒，从任务列表取数据 data = self.task_list.pop(0) # 释放锁 # 处理数据，文件路径 print(prefixed + '正在处理：',data) self.process(data) print(prefixed + '处理完成：',data) print(prefixed + '关闭消费者线程') def process(self, file_path:str): file_base_name = os.path.basename(file_path) # 文件名称 file_dir_name = os.path.dirname(file_path) # 文件路径名 file_name, file_extension = os.path.splitext(file_base_name) # 文件名 + 后缀名 final_path = os.path.join(self.tem_save_path, file_name + '.md') llm_output_path = os.path.join(self.chat_output_path, file_name + '.md') # 步骤0：构建llm推理模型，注意模型也涉及多线程,在线api可以不管 model = deepseek_online.Model_() # 步骤1：读取并处理文档 if file_extension == '.docx': # 读取器设置, 保证每个线程都有自己的处理器，避免竞态 generator = DocxReader() # docx读取器 generator.read(file_path) elif file_extension == '.pdf': # 读取器设置, 保证每个线程都有自己的处理器，避免竞态 generator = PdfReader() # pdf读取器 generator.read(file_path) else: print("未知的文件类型:", file_extension) return # 步骤2：进入文档生成器，调用模型处理 for idx, c in enumerate(generator.text_generator()): # c 就是文本内容，str llm_output = '' try: llm_output = model.chat_stream(c) except Exception as e: print("LLM推理时发生意外错误：",e) # 步骤3：保存LLM输出 if idx > 0: llm_output_path = llm_output_path.replace('.md', '-'+str(idx)+'.md') with open(llm_output_path, 'w', encoding='utf-8') as f_llm: f_llm.write(llm_output) # 步骤4：写入文件，结束文档处理，重置docx读取器状态 generator.write(final_path, 'w', 'utf-8') if __name__ == '__main__': # 测试代码 # r = Reader() # worker_iter_times = 2 # def worker(word:str): # for i in range(worker_iter_times): # print('i want to say: ,', word) # r.add_task(str(i)) # workers = 10 # threads = [threading.Thread(target=worker, kwargs={"word": str(i)}) for i in range(workers)] # for t in threads: # t.start() # for t in threads: # t.join() # print(f'now task list count is {worker_iter_times*workers}: ',r.get_task_list_len()) # del threads # part 2 # r = Reader() # r.start() # r.add_task(r'D:\code\rag_tools\终版-（0119）沭阳县循环经济产业园污废水资源化项目方案简述(3).docx') # print(r.task_list) # r.add_task(r'D:\code\rag_tools\（0119）沭阳县循环经济产业园污废水资源化项目方案简述(3).pdf') # print(r.task_list) # for i in range(10): # r.add_task('a') # r.stop() # part3测试process函数 r = Reader() # r.process(r'D:\code\rag_tools\uploads\终版-（0119）沭阳县循环经济产业园污废水资源化项目方案简述(3).docx') # r.process(r'D:\code\rag_tools\uploads\（0119）沭阳县循环经济产业园污废水资源化项目方案简述(3).pdf') r.process(r"D:\code\tem_1103\RAG资料库—开发\污水处理项目方案\TP-1-3 BioSecure生物安全性保障工艺技术包技术方案模板-V1.docx")