from typing import Tuple, Optional, Generator, List from toolbox import update_ui, update_ui_lastest_msg, get_conf import os, tarfile, requests, time, re class ArxivPaperProcessor: """Arxiv论文处理器类""" def __init__(self): self.supported_encodings = ['utf-8', 'latin1', 'gbk', 'gb2312', 'ascii'] self.arxiv_cache_dir = get_conf("ARXIV_CACHE_DIR") def download_and_extract(self, txt: str, chatbot, history) -> Generator[Optional[Tuple[str, str]], None, None]: """ Step 1: 下载和提取arxiv论文 返回: 生成器: (project_folder, arxiv_id) """ try: if txt == "": chatbot.append(("", "请输入arxiv论文链接或ID")) yield from update_ui(chatbot=chatbot, history=history) return project_folder, arxiv_id = self.arxiv_download(txt, chatbot, history) if project_folder is None or arxiv_id is None: return if not os.path.exists(project_folder): chatbot.append((txt, f"找不到项目文件夹: {project_folder}")) yield from update_ui(chatbot=chatbot, history=history) return # 期望的返回值 yield project_folder, arxiv_id except Exception as e: print(e) # yield from update_ui_lastest_msg( # "下载失败，请手动下载latex源码：请前往arxiv打开此论文下载页面，点other Formats，然后download source。", # chatbot=chatbot, history=history) return def arxiv_download(self, txt: str, chatbot, history) -> Tuple[str, str]: """ 下载arxiv论文并提取 返回: (project_folder, arxiv_id) """ def is_float(s: str) -> bool: try: float(s) return True except ValueError: return False if txt.startswith('https://arxiv.org/pdf/'): arxiv_id = txt.split('/')[-1] # 2402.14207v2.pdf txt = arxiv_id.split('v')[0] # 2402.14207 if ('.' in txt) and ('/' not in txt) and is_float(txt): # is arxiv ID txt = 'https://arxiv.org/abs/' + txt.strip() if ('.' in txt) and ('/' not in txt) and is_float(txt[:10]): # is arxiv ID txt = 'https://arxiv.org/abs/' + txt[:10] if not txt.startswith('https://arxiv.org'): chatbot.append((txt, "不是有效的arxiv链接或ID")) # yield from update_ui(chatbot=chatbot, history=history) return None, None # 返回两个值，即使其中一个为None chatbot.append([f"检测到arxiv文档连接", '尝试下载 ...']) # yield from update_ui(chatbot=chatbot, history=history) url_ = txt # https://arxiv.org/abs/1707.06690 if not txt.startswith('https://arxiv.org/abs/'): msg = f"解析arxiv网址失败, 期望格式例如: https://arxiv.org/abs/1707.06690。实际得到格式: {url_}。" # yield from update_ui_lastest_msg(msg, chatbot=chatbot, history=history) # 刷新界面 return None, None # 返回两个值，即使其中一个为None arxiv_id = url_.split('/')[-1].split('v')[0] dst = os.path.join(self.arxiv_cache_dir, arxiv_id, f'{arxiv_id}.tar.gz') project_folder = os.path.join(self.arxiv_cache_dir, arxiv_id) success = self.download_arxiv_paper(url_, dst, chatbot, history) # if os.path.exists(dst) and get_conf('allow_cache'): # # yield from update_ui_lastest_msg(f"调用缓存 {arxiv_id}", chatbot=chatbot, history=history) # 刷新界面 # success = True # else: # # yield from update_ui_lastest_msg(f"开始下载 {arxiv_id}", chatbot=chatbot, history=history) # 刷新界面 # success = self.download_arxiv_paper(url_, dst, chatbot, history) # # yield from update_ui_lastest_msg(f"下载完成 {arxiv_id}", chatbot=chatbot, history=history) # 刷新界面 if not success: # chatbot.append([f"下载失败 {arxiv_id}", ""]) # yield from update_ui(chatbot=chatbot, history=history) raise tarfile.ReadError(f"论文下载失败 {arxiv_id}") # yield from update_ui_lastest_msg(f"开始解压 {arxiv_id}", chatbot=chatbot, history=history) # 刷新界面 extract_dst = self.extract_tar_file(dst, project_folder, chatbot, history) # yield from update_ui_lastest_msg(f"解压完成 {arxiv_id}", chatbot=chatbot, history=history) # 刷新界面 return extract_dst, arxiv_id def download_arxiv_paper(self, url_: str, dst: str, chatbot, history) -> bool: """下载arxiv论文""" try: proxies = get_conf('proxies') for url_tar in [url_.replace('/abs/', '/src/'), url_.replace('/abs/', '/e-print/')]: r = requests.get(url_tar, proxies=proxies) if r.status_code == 200: with open(dst, 'wb+') as f: f.write(r.content) return True return False except requests.RequestException as e: # chatbot.append((f"下载失败 {url_}", str(e))) # yield from update_ui(chatbot=chatbot, history=history) return False def extract_tar_file(self, file_path: str, dest_dir: str, chatbot, history) -> str: """解压arxiv论文""" try: with tarfile.open(file_path, 'r:gz') as tar: tar.extractall(path=dest_dir) return dest_dir except tarfile.ReadError as e: chatbot.append((f"解压失败 {file_path}", str(e))) yield from update_ui(chatbot=chatbot, history=history) raise e def find_main_tex_file(self, tex_files: list) -> str: """查找主TEX文件""" for tex_file in tex_files: with open(tex_file, 'r', encoding='utf-8', errors='ignore') as f: content = f.read() if r'\documentclass' in content: return tex_file return max(tex_files, key=lambda x: os.path.getsize(x)) def read_file_with_encoding(self, file_path: str) -> Optional[str]: """使用多种编码尝试读取文件""" for encoding in self.supported_encodings: try: with open(file_path, 'r', encoding=encoding) as f: return f.read() except UnicodeDecodeError: continue return None def process_tex_content(self, content: str, base_path: str, processed_files=None) -> str: """处理TEX内容，包括递归处理包含的文件""" if processed_files is None: processed_files = set() include_patterns = [ r'\\input{([^}]+)}', r'\\include{([^}]+)}', r'\\subfile{([^}]+)}', r'\\input\s+([^\s{]+)', ] for pattern in include_patterns: matches = re.finditer(pattern, content) for match in matches: include_file = match.group(1) if not include_file.endswith('.tex'): include_file += '.tex' include_path = os.path.join(base_path, include_file) include_path = os.path.normpath(include_path) if include_path in processed_files: continue processed_files.add(include_path) if os.path.exists(include_path): included_content = self.read_file_with_encoding(include_path) if included_content: included_content = self.process_tex_content( included_content, os.path.dirname(include_path), processed_files ) content = content.replace(match.group(0), included_content) return content def merge_tex_files(self, folder_path: str, chatbot, history) -> Optional[str]: """ Step 2: 合并TEX文件 返回: 合并后的内容 """ try: tex_files = [] for root, _, files in os.walk(folder_path): tex_files.extend([os.path.join(root, f) for f in files if f.endswith('.tex')]) if not tex_files: chatbot.append(("", "未找到任何TEX文件")) yield from update_ui(chatbot=chatbot, history=history) return None main_tex_file = self.find_main_tex_file(tex_files) chatbot.append(("", f"找到主TEX文件：{os.path.basename(main_tex_file)}")) yield from update_ui(chatbot=chatbot, history=history) tex_content = self.read_file_with_encoding(main_tex_file) if tex_content is None: chatbot.append(("", "无法读取TEX文件，可能是编码问题")) yield from update_ui(chatbot=chatbot, history=history) return None full_content = self.process_tex_content( tex_content, os.path.dirname(main_tex_file) ) cleaned_content = self.clean_tex_content(full_content) chatbot.append(("", f"成功处理所有TEX文件：\n" f"- 原始内容大小：{len(full_content)}字符\n" f"- 清理后内容大小：{len(cleaned_content)}字符" )) yield from update_ui(chatbot=chatbot, history=history) # 添加标题和摘要提取 title = "" abstract = "" if tex_content: # 提取标题 title_match = re.search(r'\\title{([^}]*)}', tex_content) if title_match: title = title_match.group(1) # 提取摘要 abstract_match = re.search(r'\\begin{abstract}(.*?)\\end{abstract}', tex_content, re.DOTALL) if abstract_match: abstract = abstract_match.group(1) # 按token限制分段 def split_by_token_limit(text: str, token_limit: int = 1024) -> List[str]: segments = [] current_segment = [] current_tokens = 0 for line in text.split('\n'): line_tokens = len(line.split()) if current_tokens + line_tokens > token_limit: segments.append('\n'.join(current_segment)) current_segment = [line] current_tokens = line_tokens else: current_segment.append(line) current_tokens += line_tokens if current_segment: segments.append('\n'.join(current_segment)) return segments text_segments = split_by_token_limit(cleaned_content) return { 'title': title, 'abstract': abstract, 'segments': text_segments } except Exception as e: chatbot.append(("", f"处理TEX文件时发生错误：{str(e)}")) yield from update_ui(chatbot=chatbot, history=history) return None @staticmethod def clean_tex_content(content: str) -> str: """清理TEX内容""" content = re.sub(r'(?m)%.*$', '', content) # 移除注释 content = re.sub(r'\\cite{[^}]*}', '', content) # 移除引用 content = re.sub(r'\\label{[^}]*}', '', content) # 移除标签 content = re.sub(r'\s+', ' ', content) # 规范化空白 return content.strip() if __name__ == "__main__": # 测试 arxiv_download 函数 processor = ArxivPaperProcessor() chatbot = [] history = [] # 测试不同格式的输入 test_inputs = [ "https://arxiv.org/abs/2402.14207", # 标准格式 "https://arxiv.org/pdf/2402.14207.pdf", # PDF链接格式 "2402.14207", # 纯ID格式 "2402.14207v1", # 带版本号的ID格式 "https://invalid.url", # 无效URL测试 ] for input_url in test_inputs: print(f"\n测试输入: {input_url}") try: project_folder, arxiv_id = processor.arxiv_download(input_url, chatbot, history) if project_folder and arxiv_id: print(f"下载成功:") print(f"- 项目文件夹: {project_folder}") print(f"- Arxiv ID: {arxiv_id}") print(f"- 文件夹是否存在: {os.path.exists(project_folder)}") else: print("下载失败: 返回值为 None") except Exception as e: print(f"发生错误: {str(e)}")