镜像自地址
https://github.com/binary-husky/gpt_academic.git
已同步 2025-12-06 14:36:48 +00:00
511 行
18 KiB
Python
511 行
18 KiB
Python
import os
|
|
import re
|
|
import requests
|
|
import tarfile
|
|
import logging
|
|
from dataclasses import dataclass
|
|
from typing import Generator, List, Tuple, Optional, Dict, Set
|
|
from pathlib import Path
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
|
|
|
|
@dataclass
|
|
class ArxivFragment:
|
|
"""Arxiv论文片段数据类"""
|
|
file_path: str
|
|
content: str
|
|
segment_index: int
|
|
total_segments: int
|
|
rel_path: str
|
|
segment_type: str
|
|
title: str
|
|
abstract: str
|
|
section: str
|
|
is_appendix: bool
|
|
|
|
|
|
class SmartArxivSplitter:
|
|
def __init__(self,
|
|
char_range: Tuple[int, int],
|
|
root_dir: str = "gpt_log/arxiv_cache",
|
|
proxies: Optional[Dict[str, str]] = None,
|
|
max_workers: int = 4):
|
|
|
|
self.min_chars, self.max_chars = char_range
|
|
self.root_dir = Path(root_dir)
|
|
self.root_dir.mkdir(parents=True, exist_ok=True)
|
|
self.proxies = proxies or {}
|
|
self.max_workers = max_workers
|
|
|
|
# 定义特殊环境模式
|
|
self._init_patterns()
|
|
|
|
# 配置日志
|
|
logging.basicConfig(level=logging.INFO,
|
|
format='%(asctime)s - %(levelname)s - %(message)s')
|
|
|
|
def _init_patterns(self):
|
|
"""初始化LaTeX环境和命令模式"""
|
|
self.special_envs = {
|
|
'math': [r'\\begin{(equation|align|gather|eqnarray)\*?}.*?\\end{\1\*?}',
|
|
r'\$\$.*?\$\$', r'\$[^$]+\$'],
|
|
'table': [r'\\begin{(table|tabular)\*?}.*?\\end{\1\*?}'],
|
|
'figure': [r'\\begin{figure\*?}.*?\\end{figure\*?}'],
|
|
'algorithm': [r'\\begin{(algorithm|algorithmic)}.*?\\end{\1}']
|
|
}
|
|
|
|
self.section_patterns = [
|
|
r'\\(sub)*section\{([^}]+)\}',
|
|
r'\\chapter\{([^}]+)\}'
|
|
]
|
|
|
|
self.include_patterns = [
|
|
r'\\(input|include|subfile)\{([^}]+)\}'
|
|
]
|
|
|
|
def _find_main_tex_file(self, directory: str) -> Optional[str]:
|
|
"""查找主TEX文件"""
|
|
tex_files = list(Path(directory).rglob("*.tex"))
|
|
if not tex_files:
|
|
return None
|
|
|
|
# 按以下优先级查找：
|
|
# 1. 包含documentclass的文件
|
|
# 2. 文件名为main.tex
|
|
# 3. 最大的tex文件
|
|
for tex_file in tex_files:
|
|
try:
|
|
content = self._read_file(str(tex_file))
|
|
if content and r'\documentclass' in content:
|
|
return str(tex_file)
|
|
if tex_file.name.lower() == 'main.tex':
|
|
return str(tex_file)
|
|
except Exception:
|
|
continue
|
|
|
|
return str(max(tex_files, key=lambda x: x.stat().st_size))
|
|
|
|
def _resolve_includes(self, tex_file: str, processed: Set[str] = None) -> List[str]:
|
|
"""递归解析tex文件中的include/input命令"""
|
|
if processed is None:
|
|
processed = set()
|
|
|
|
if tex_file in processed:
|
|
return []
|
|
|
|
processed.add(tex_file)
|
|
result = [tex_file]
|
|
content = self._read_file(tex_file)
|
|
|
|
if not content:
|
|
return result
|
|
|
|
base_dir = Path(tex_file).parent
|
|
for pattern in self.include_patterns:
|
|
for match in re.finditer(pattern, content):
|
|
included_file = match.group(2)
|
|
if not included_file.endswith('.tex'):
|
|
included_file += '.tex'
|
|
|
|
# 构建完整路径
|
|
full_path = str(base_dir / included_file)
|
|
if os.path.exists(full_path) and full_path not in processed:
|
|
result.extend(self._resolve_includes(full_path, processed))
|
|
|
|
return result
|
|
|
|
def _smart_split(self, content: str) -> List[Tuple[str, str, bool]]:
|
|
"""智能分割TEX内容，确保在字符范围内并保持语义完整性"""
|
|
content = self._preprocess_content(content)
|
|
segments = []
|
|
current_buffer = []
|
|
current_length = 0
|
|
current_section = "Unknown Section"
|
|
is_appendix = False
|
|
|
|
# 保护特殊环境
|
|
protected_blocks = {}
|
|
content = self._protect_special_environments(content, protected_blocks)
|
|
|
|
# 按段落分割
|
|
paragraphs = re.split(r'\n\s*\n', content)
|
|
|
|
for para in paragraphs:
|
|
para = para.strip()
|
|
if not para:
|
|
continue
|
|
|
|
# 恢复特殊环境
|
|
para = self._restore_special_environments(para, protected_blocks)
|
|
|
|
# 更新章节信息
|
|
section_info = self._get_section_info(para, content)
|
|
if section_info:
|
|
current_section, is_appendix = section_info
|
|
|
|
# 判断是否是特殊环境
|
|
if self._is_special_environment(para):
|
|
# 处理当前缓冲区
|
|
if current_buffer:
|
|
segments.append((
|
|
'\n'.join(current_buffer),
|
|
current_section,
|
|
is_appendix
|
|
))
|
|
current_buffer = []
|
|
current_length = 0
|
|
|
|
# 添加特殊环境作为独立片段
|
|
segments.append((para, current_section, is_appendix))
|
|
continue
|
|
|
|
# 处理普通段落
|
|
sentences = self._split_into_sentences(para)
|
|
for sentence in sentences:
|
|
sentence = sentence.strip()
|
|
if not sentence:
|
|
continue
|
|
|
|
sent_length = len(sentence)
|
|
new_length = current_length + sent_length + (1 if current_buffer else 0)
|
|
|
|
if new_length <= self.max_chars:
|
|
current_buffer.append(sentence)
|
|
current_length = new_length
|
|
else:
|
|
# 如果当前缓冲区达到最小长度要求
|
|
if current_length >= self.min_chars:
|
|
segments.append((
|
|
'\n'.join(current_buffer),
|
|
current_section,
|
|
is_appendix
|
|
))
|
|
current_buffer = [sentence]
|
|
current_length = sent_length
|
|
else:
|
|
# 尝试将过长的句子分割
|
|
split_sentences = self._split_long_sentence(sentence)
|
|
for split_sent in split_sentences:
|
|
if current_length + len(split_sent) <= self.max_chars:
|
|
current_buffer.append(split_sent)
|
|
current_length += len(split_sent) + 1
|
|
else:
|
|
segments.append((
|
|
'\n'.join(current_buffer),
|
|
current_section,
|
|
is_appendix
|
|
))
|
|
current_buffer = [split_sent]
|
|
current_length = len(split_sent)
|
|
|
|
# 处理剩余的缓冲区
|
|
if current_buffer:
|
|
segments.append((
|
|
'\n'.join(current_buffer),
|
|
current_section,
|
|
is_appendix
|
|
))
|
|
|
|
return segments
|
|
|
|
def _split_into_sentences(self, text: str) -> List[str]:
|
|
"""将文本分割成句子"""
|
|
return re.split(r'(?<=[.!?。！？])\s+', text)
|
|
|
|
def _split_long_sentence(self, sentence: str) -> List[str]:
|
|
"""智能分割过长的句子"""
|
|
if len(sentence) <= self.max_chars:
|
|
return [sentence]
|
|
|
|
result = []
|
|
while sentence:
|
|
# 在最大长度位置寻找合适的分割点
|
|
split_pos = self._find_split_position(sentence[:self.max_chars])
|
|
if split_pos <= 0:
|
|
split_pos = self.max_chars
|
|
|
|
result.append(sentence[:split_pos])
|
|
sentence = sentence[split_pos:].strip()
|
|
|
|
return result
|
|
|
|
def _find_split_position(self, text: str) -> int:
|
|
"""找到合适的句子分割位置"""
|
|
# 优先在标点符号处分割
|
|
punctuation_match = re.search(r'[,，;；]\s*', text[::-1])
|
|
if punctuation_match:
|
|
return len(text) - punctuation_match.end()
|
|
|
|
# 其次在空白字符处分割
|
|
space_match = re.search(r'\s+', text[::-1])
|
|
if space_match:
|
|
return len(text) - space_match.end()
|
|
|
|
return -1
|
|
|
|
def _protect_special_environments(self, content: str, protected_blocks: Dict[str, str]) -> str:
|
|
"""保护特殊环境内容"""
|
|
for env_type, patterns in self.special_envs.items():
|
|
for pattern in patterns:
|
|
content = re.sub(
|
|
pattern,
|
|
lambda m: self._store_protected_block(m.group(0), protected_blocks),
|
|
content,
|
|
flags=re.DOTALL
|
|
)
|
|
return content
|
|
|
|
def _store_protected_block(self, content: str, protected_blocks: Dict[str, str]) -> str:
|
|
"""存储受保护的内容块"""
|
|
placeholder = f"PROTECTED_{len(protected_blocks)}"
|
|
protected_blocks[placeholder] = content
|
|
return placeholder
|
|
|
|
def _restore_special_environments(self, content: str, protected_blocks: Dict[str, str]) -> str:
|
|
"""恢复特殊环境内容"""
|
|
for placeholder, original in protected_blocks.items():
|
|
content = content.replace(placeholder, original)
|
|
return content
|
|
|
|
def _is_special_environment(self, text: str) -> bool:
|
|
"""判断是否是特殊环境"""
|
|
for patterns in self.special_envs.values():
|
|
for pattern in patterns:
|
|
if re.search(pattern, text, re.DOTALL):
|
|
return True
|
|
return False
|
|
|
|
def _preprocess_content(self, content: str) -> str:
|
|
"""预处理TEX内容"""
|
|
# 移除注释
|
|
content = re.sub(r'(?m)%.*$', '', content)
|
|
# 规范化空白字符
|
|
content = re.sub(r'\s+', ' ', content)
|
|
content = re.sub(r'\n\s*\n', '\n\n', content)
|
|
# 移除不必要的命令
|
|
content = re.sub(r'\\(label|ref|cite)\{[^}]*\}', '', content)
|
|
return content.strip()
|
|
|
|
def process(self, arxiv_id_or_url: str) -> Generator[ArxivFragment, None, None]:
|
|
"""处理单篇arxiv论文"""
|
|
try:
|
|
arxiv_id = self._normalize_arxiv_id(arxiv_id_or_url)
|
|
paper_dir = self._download_and_extract(arxiv_id)
|
|
|
|
# 查找主tex文件
|
|
main_tex = self._find_main_tex_file(paper_dir)
|
|
if not main_tex:
|
|
raise RuntimeError(f"No main tex file found in {paper_dir}")
|
|
|
|
# 获取所有相关tex文件
|
|
tex_files = self._resolve_includes(main_tex)
|
|
|
|
# 处理所有tex文件
|
|
fragments = []
|
|
with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
|
|
future_to_file = {
|
|
executor.submit(self._process_single_tex, file_path): file_path
|
|
for file_path in tex_files
|
|
}
|
|
|
|
for future in as_completed(future_to_file):
|
|
try:
|
|
fragments.extend(future.result())
|
|
except Exception as e:
|
|
logging.error(f"Error processing file: {e}")
|
|
|
|
# 重新计算片段索引
|
|
fragments.sort(key=lambda x: (x.rel_path, x.segment_index))
|
|
total_fragments = len(fragments)
|
|
|
|
for i, fragment in enumerate(fragments):
|
|
fragment.segment_index = i
|
|
fragment.total_segments = total_fragments
|
|
yield fragment
|
|
|
|
except Exception as e:
|
|
logging.error(f"Error processing paper {arxiv_id_or_url}: {e}")
|
|
raise RuntimeError(f"Failed to process paper: {str(e)}")
|
|
|
|
def _normalize_arxiv_id(self, input_str: str) -> str:
|
|
"""规范化arxiv ID"""
|
|
if input_str.startswith('https://arxiv.org/'):
|
|
if '/pdf/' in input_str:
|
|
return input_str.split('/pdf/')[-1].split('v')[0]
|
|
return input_str.split('/abs/')[-1].split('v')[0]
|
|
return input_str.split('v')[0]
|
|
|
|
def _download_and_extract(self, arxiv_id: str) -> str:
|
|
"""下载并解压arxiv论文源码"""
|
|
paper_dir = self.root_dir / arxiv_id
|
|
tar_path = paper_dir / f"{arxiv_id}.tar.gz"
|
|
|
|
# 检查缓存
|
|
if paper_dir.exists() and any(paper_dir.iterdir()):
|
|
logging.info(f"Using cached version for {arxiv_id}")
|
|
return str(paper_dir)
|
|
|
|
paper_dir.mkdir(exist_ok=True)
|
|
|
|
urls = [
|
|
f"https://arxiv.org/src/{arxiv_id}",
|
|
f"https://arxiv.org/e-print/{arxiv_id}"
|
|
]
|
|
|
|
for url in urls:
|
|
try:
|
|
logging.info(f"Downloading from {url}")
|
|
response = requests.get(url, proxies=self.proxies)
|
|
if response.status_code == 200:
|
|
tar_path.write_bytes(response.content)
|
|
with tarfile.open(tar_path, 'r:gz') as tar:
|
|
tar.extractall(path=paper_dir)
|
|
return str(paper_dir)
|
|
except Exception as e:
|
|
logging.warning(f"Download failed for {url}: {e}")
|
|
continue
|
|
|
|
raise RuntimeError(f"Failed to download paper {arxiv_id}")
|
|
|
|
def _read_file(self, file_path: str) -> Optional[str]:
|
|
"""使用多种编码尝试读取文件"""
|
|
encodings = ['utf-8', 'latin1', 'gbk', 'gb2312', 'ascii']
|
|
for encoding in encodings:
|
|
try:
|
|
with open(file_path, 'r', encoding=encoding) as f:
|
|
return f.read()
|
|
except UnicodeDecodeError:
|
|
continue
|
|
logging.warning(f"Failed to read file {file_path} with all encodings")
|
|
return None
|
|
|
|
def _extract_metadata(self, content: str) -> Tuple[str, str]:
|
|
"""提取论文标题和摘要"""
|
|
title = ""
|
|
abstract = ""
|
|
|
|
# 提取标题
|
|
title_patterns = [
|
|
r'\\title{([^}]*)}',
|
|
r'\\Title{([^}]*)}'
|
|
]
|
|
for pattern in title_patterns:
|
|
match = re.search(pattern, content)
|
|
if match:
|
|
title = match.group(1)
|
|
title = re.sub(r'\\[a-zA-Z]+{([^}]*)}', r'\1', title)
|
|
break
|
|
|
|
# 提取摘要
|
|
abstract_patterns = [
|
|
r'\\begin{abstract}(.*?)\\end{abstract}',
|
|
r'\\abstract{([^}]*)}'
|
|
]
|
|
for pattern in abstract_patterns:
|
|
match = re.search(pattern, content, re.DOTALL)
|
|
if match:
|
|
abstract = match.group(1).strip()
|
|
abstract = re.sub(r'\\[a-zA-Z]+{([^}]*)}', r'\1', abstract)
|
|
break
|
|
|
|
return title.strip(), abstract.strip()
|
|
|
|
def _get_section_info(self, para: str, content: str) -> Optional[Tuple[str, bool]]:
|
|
"""获取段落所属的章节信息"""
|
|
section = "Unknown Section"
|
|
is_appendix = False
|
|
|
|
# 查找所有章节标记
|
|
all_sections = []
|
|
for pattern in self.section_patterns:
|
|
for match in re.finditer(pattern, content):
|
|
all_sections.append((match.start(), match.group(2)))
|
|
|
|
# 查找appendix标记
|
|
appendix_pos = content.find(r'\appendix')
|
|
|
|
# 确定当前章节
|
|
para_pos = content.find(para)
|
|
if para_pos >= 0:
|
|
current_section = None
|
|
for sec_pos, sec_title in sorted(all_sections):
|
|
if sec_pos > para_pos:
|
|
break
|
|
current_section = sec_title
|
|
|
|
if current_section:
|
|
section = current_section
|
|
is_appendix = appendix_pos >= 0 and para_pos > appendix_pos
|
|
|
|
return section, is_appendix
|
|
|
|
return None
|
|
|
|
def _process_single_tex(self, file_path: str) -> List[ArxivFragment]:
|
|
"""处理单个TEX文件"""
|
|
try:
|
|
content = self._read_file(file_path)
|
|
if not content:
|
|
return []
|
|
|
|
# 提取元数据
|
|
is_main = r'\documentclass' in content
|
|
title = ""
|
|
abstract = ""
|
|
if is_main:
|
|
title, abstract = self._extract_metadata(content)
|
|
|
|
# 智能分割内容
|
|
segments = self._smart_split(content)
|
|
fragments = []
|
|
|
|
for i, (segment_content, section, is_appendix) in enumerate(segments):
|
|
if segment_content.strip():
|
|
segment_type = 'text'
|
|
for env_type, patterns in self.special_envs.items():
|
|
if any(re.search(pattern, segment_content, re.DOTALL)
|
|
for pattern in patterns):
|
|
segment_type = env_type
|
|
break
|
|
|
|
fragments.append(ArxivFragment(
|
|
file_path=file_path,
|
|
content=segment_content,
|
|
segment_index=i,
|
|
total_segments=len(segments),
|
|
rel_path=os.path.relpath(file_path, str(self.root_dir)),
|
|
segment_type=segment_type,
|
|
title=title,
|
|
abstract=abstract,
|
|
section=section,
|
|
is_appendix=is_appendix
|
|
))
|
|
|
|
return fragments
|
|
|
|
except Exception as e:
|
|
logging.error(f"Error processing file {file_path}: {e}")
|
|
return []
|
|
|
|
def main():
|
|
"""使用示例"""
|
|
# 创建分割器实例
|
|
splitter = SmartArxivSplitter(
|
|
char_range=(1000, 1200),
|
|
root_dir="gpt_log/arxiv_cache"
|
|
)
|
|
|
|
# 处理论文
|
|
for fragment in splitter.process("2411.03663"):
|
|
print(f"Segment {fragment.segment_index + 1}/{fragment.total_segments}")
|
|
print(f"Length: {len(fragment.content)}")
|
|
print(f"Section: {fragment.section}")
|
|
print(f"Title: {fragment.file_path}")
|
|
|
|
print(fragment.content)
|
|
print("-" * 80)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|