123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562 |
- import os, shutil
- import re
- import numpy as np
- PRESERVE = 0
- TRANSFORM = 1
- pj = os.path.join
- class LinkedListNode():
- """
- Linked List Node
- """
- def __init__(self, string, preserve=True) -> None:
- self.string = string
- self.preserve = preserve
- self.next = None
- self.range = None
- # self.begin_line = 0
- # self.begin_char = 0
- def convert_to_linklist(text, mask):
- root = LinkedListNode("", preserve=True)
- current_node = root
- for c, m, i in zip(text, mask, range(len(text))):
- if (m==PRESERVE and current_node.preserve) \
- or (m==TRANSFORM and not current_node.preserve):
- # add
- current_node.string += c
- else:
- current_node.next = LinkedListNode(c, preserve=(m==PRESERVE))
- current_node = current_node.next
- return root
- def post_process(root):
- # 修复括号
- node = root
- while True:
- string = node.string
- if node.preserve:
- node = node.next
- if node is None: break
- continue
- def break_check(string):
- str_stack = [""] # (lv, index)
- for i, c in enumerate(string):
- if c == '{':
- str_stack.append('{')
- elif c == '}':
- if len(str_stack) == 1:
- print('stack fix')
- return i
- str_stack.pop(-1)
- else:
- str_stack[-1] += c
- return -1
- bp = break_check(string)
- if bp == -1:
- pass
- elif bp == 0:
- node.string = string[:1]
- q = LinkedListNode(string[1:], False)
- q.next = node.next
- node.next = q
- else:
- node.string = string[:bp]
- q = LinkedListNode(string[bp:], False)
- q.next = node.next
- node.next = q
- node = node.next
- if node is None: break
- # 屏蔽空行和太短的句子
- node = root
- while True:
- if len(node.string.strip('\n').strip(''))==0: node.preserve = True
- if len(node.string.strip('\n').strip(''))<42: node.preserve = True
- node = node.next
- if node is None: break
- node = root
- while True:
- if node.next and node.preserve and node.next.preserve:
- node.string += node.next.string
- node.next = node.next.next
- node = node.next
- if node is None: break
- # 将前后断行符脱离
- node = root
- prev_node = None
- while True:
- if not node.preserve:
- lstriped_ = node.string.lstrip().lstrip('\n')
- if (prev_node is not None) and (prev_node.preserve) and (len(lstriped_)!=len(node.string)):
- prev_node.string += node.string[:-len(lstriped_)]
- node.string = lstriped_
- rstriped_ = node.string.rstrip().rstrip('\n')
- if (node.next is not None) and (node.next.preserve) and (len(rstriped_)!=len(node.string)):
- node.next.string = node.string[len(rstriped_):] + node.next.string
- node.string = rstriped_
- # =====
- prev_node = node
- node = node.next
- if node is None: break
- # 标注节点的行数范围
- node = root
- n_line = 0
- expansion = 2
- while True:
- n_l = node.string.count('\n')
- node.range = [n_line-expansion, n_line+n_l+expansion] # 失败时,扭转的范围
- n_line = n_line+n_l
- node = node.next
- if node is None: break
- return root
- """
- =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
- Latex segmentation with a binary mask (PRESERVE=0, TRANSFORM=1)
- =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
- """
- def set_forbidden_text(text, mask, pattern, flags=0):
- """
- Add a preserve text area in this paper
- e.g. with pattern = r"\\begin\{algorithm\}(.*?)\\end\{algorithm\}"
- you can mask out (mask = PRESERVE so that text become untouchable for GPT)
- everything between "\begin{equation}" and "\end{equation}"
- """
- if isinstance(pattern, list): pattern = '|'.join(pattern)
- pattern_compile = re.compile(pattern, flags)
- for res in pattern_compile.finditer(text):
- mask[res.span()[0]:res.span()[1]] = PRESERVE
- return text, mask
- def reverse_forbidden_text(text, mask, pattern, flags=0, forbid_wrapper=True):
- """
- Move area out of preserve area (make text editable for GPT)
- count the number of the braces so as to catch compelete text area.
- e.g.
- \begin{abstract} blablablablablabla. \end{abstract}
- """
- if isinstance(pattern, list): pattern = '|'.join(pattern)
- pattern_compile = re.compile(pattern, flags)
- for res in pattern_compile.finditer(text):
- if not forbid_wrapper:
- mask[res.span()[0]:res.span()[1]] = TRANSFORM
- else:
- mask[res.regs[0][0]: res.regs[1][0]] = PRESERVE # '\\begin{abstract}'
- mask[res.regs[1][0]: res.regs[1][1]] = TRANSFORM # abstract
- mask[res.regs[1][1]: res.regs[0][1]] = PRESERVE # abstract
- return text, mask
- def set_forbidden_text_careful_brace(text, mask, pattern, flags=0):
- """
- Add a preserve text area in this paper (text become untouchable for GPT).
- count the number of the braces so as to catch compelete text area.
- e.g.
- \caption{blablablablabla\texbf{blablabla}blablabla.}
- """
- pattern_compile = re.compile(pattern, flags)
- for res in pattern_compile.finditer(text):
- brace_level = -1
- p = begin = end = res.regs[0][0]
- for _ in range(1024*16):
- if text[p] == '}' and brace_level == 0: break
- elif text[p] == '}': brace_level -= 1
- elif text[p] == '{': brace_level += 1
- p += 1
- end = p+1
- mask[begin:end] = PRESERVE
- return text, mask
- def reverse_forbidden_text_careful_brace(text, mask, pattern, flags=0, forbid_wrapper=True):
- """
- Move area out of preserve area (make text editable for GPT)
- count the number of the braces so as to catch compelete text area.
- e.g.
- \caption{blablablablabla\texbf{blablabla}blablabla.}
- """
- pattern_compile = re.compile(pattern, flags)
- for res in pattern_compile.finditer(text):
- brace_level = 0
- p = begin = end = res.regs[1][0]
- for _ in range(1024*16):
- if text[p] == '}' and brace_level == 0: break
- elif text[p] == '}': brace_level -= 1
- elif text[p] == '{': brace_level += 1
- p += 1
- end = p
- mask[begin:end] = TRANSFORM
- if forbid_wrapper:
- mask[res.regs[0][0]:begin] = PRESERVE
- mask[end:res.regs[0][1]] = PRESERVE
- return text, mask
- def set_forbidden_text_begin_end(text, mask, pattern, flags=0, limit_n_lines=42):
- """
- Find all \begin{} ... \end{} text block that with less than limit_n_lines lines.
- Add it to preserve area
- """
- pattern_compile = re.compile(pattern, flags)
- def search_with_line_limit(text, mask):
- for res in pattern_compile.finditer(text):
- cmd = res.group(1) # begin{what}
- this = res.group(2) # content between begin and end
- this_mask = mask[res.regs[2][0]:res.regs[2][1]]
- white_list = ['document', 'abstract', 'lemma', 'definition', 'sproof',
- 'em', 'emph', 'textit', 'textbf', 'itemize', 'enumerate']
- if (cmd in white_list) or this.count('\n') >= limit_n_lines: # use a magical number 42
- this, this_mask = search_with_line_limit(this, this_mask)
- mask[res.regs[2][0]:res.regs[2][1]] = this_mask
- else:
- mask[res.regs[0][0]:res.regs[0][1]] = PRESERVE
- return text, mask
- return search_with_line_limit(text, mask)
- """
- =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
- Latex Merge File
- =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
- """
- def find_main_tex_file(file_manifest, mode):
- """
- 在多Tex文档中,寻找主文件,必须包含documentclass,返回找到的第一个。
- P.S. 但愿没人把latex模板放在里面传进来 (6.25 加入判定latex模板的代码)
- """
- canidates = []
- for texf in file_manifest:
- if os.path.basename(texf).startswith('merge'):
- continue
- with open(texf, 'r', encoding='utf8', errors='ignore') as f:
- file_content = f.read()
- if r'\documentclass' in file_content:
- canidates.append(texf)
- else:
- continue
- if len(canidates) == 0:
- raise RuntimeError('无法找到一个主Tex文件(包含documentclass关键字)')
- elif len(canidates) == 1:
- return canidates[0]
- else: # if len(canidates) >= 2 通过一些Latex模板中常见(但通常不会出现在正文)的单词,对不同latex源文件扣分,取评分最高者返回
- canidates_score = []
- # 给出一些判定模板文档的词作为扣分项
- unexpected_words = ['\\LaTeX', 'manuscript', 'Guidelines', 'font', 'citations', 'rejected', 'blind review', 'reviewers']
- expected_words = ['\\input', '\\ref', '\\cite']
- for texf in canidates:
- canidates_score.append(0)
- with open(texf, 'r', encoding='utf8', errors='ignore') as f:
- file_content = f.read()
- file_content = rm_comments(file_content)
- for uw in unexpected_words:
- if uw in file_content:
- canidates_score[-1] -= 1
- for uw in expected_words:
- if uw in file_content:
- canidates_score[-1] += 1
- select = np.argmax(canidates_score) # 取评分最高者返回
- return canidates[select]
-
- def rm_comments(main_file):
- new_file_remove_comment_lines = []
- for l in main_file.splitlines():
- # 删除整行的空注释
- if l.lstrip().startswith("%"):
- pass
- else:
- new_file_remove_comment_lines.append(l)
- main_file = '\n'.join(new_file_remove_comment_lines)
- # main_file = re.sub(r"\\include{(.*?)}", r"\\input{\1}", main_file) # 将 \include 命令转换为 \input 命令
- main_file = re.sub(r'(?<!\\)%.*', '', main_file) # 使用正则表达式查找半行注释, 并替换为空字符串
- return main_file
- def find_tex_file_ignore_case(fp):
- dir_name = os.path.dirname(fp)
- base_name = os.path.basename(fp)
- # 如果输入的文件路径是正确的
- if os.path.isfile(pj(dir_name, base_name)): return pj(dir_name, base_name)
- # 如果不正确,试着加上.tex后缀试试
- if not base_name.endswith('.tex'): base_name+='.tex'
- if os.path.isfile(pj(dir_name, base_name)): return pj(dir_name, base_name)
- # 如果还找不到,解除大小写限制,再试一次
- import glob
- for f in glob.glob(dir_name+'/*.tex'):
- base_name_s = os.path.basename(fp)
- base_name_f = os.path.basename(f)
- if base_name_s.lower() == base_name_f.lower(): return f
- # 试着加上.tex后缀试试
- if not base_name_s.endswith('.tex'): base_name_s+='.tex'
- if base_name_s.lower() == base_name_f.lower(): return f
- return None
- def merge_tex_files_(project_foler, main_file, mode):
- """
- Merge Tex project recrusively
- """
- main_file = rm_comments(main_file)
- for s in reversed([q for q in re.finditer(r"\\input\{(.*?)\}", main_file, re.M)]):
- f = s.group(1)
- fp = os.path.join(project_foler, f)
- fp_ = find_tex_file_ignore_case(fp)
- if fp_:
- try:
- with open(fp_, 'r', encoding='utf-8', errors='replace') as fx: c = fx.read()
- except:
- c = f"\n\nWarning from GPT-Academic: LaTex source file is missing!\n\n"
- else:
- raise RuntimeError(f'找不到{fp},Tex源文件缺失!')
- c = merge_tex_files_(project_foler, c, mode)
- main_file = main_file[:s.span()[0]] + c + main_file[s.span()[1]:]
- return main_file
- def find_title_and_abs(main_file):
- def extract_abstract_1(text):
- pattern = r"\\abstract\{(.*?)\}"
- match = re.search(pattern, text, re.DOTALL)
- if match:
- return match.group(1)
- else:
- return None
- def extract_abstract_2(text):
- pattern = r"\\begin\{abstract\}(.*?)\\end\{abstract\}"
- match = re.search(pattern, text, re.DOTALL)
- if match:
- return match.group(1)
- else:
- return None
- def extract_title(string):
- pattern = r"\\title\{(.*?)\}"
- match = re.search(pattern, string, re.DOTALL)
- if match:
- return match.group(1)
- else:
- return None
- abstract = extract_abstract_1(main_file)
- if abstract is None:
- abstract = extract_abstract_2(main_file)
- title = extract_title(main_file)
- return title, abstract
- def merge_tex_files(project_foler, main_file, mode):
- """
- Merge Tex project recrusively
- P.S. 顺便把CTEX塞进去以支持中文
- P.S. 顺便把Latex的注释去除
- """
- main_file = merge_tex_files_(project_foler, main_file, mode)
- main_file = rm_comments(main_file)
- if mode == 'translate_zh':
- # find paper documentclass
- pattern = re.compile(r'\\documentclass.*\n')
- match = pattern.search(main_file)
- assert match is not None, "Cannot find documentclass statement!"
- position = match.end()
- add_ctex = '\\usepackage{ctex}\n'
- add_url = '\\usepackage{url}\n' if '{url}' not in main_file else ''
- main_file = main_file[:position] + add_ctex + add_url + main_file[position:]
- # fontset=windows
- import platform
- main_file = re.sub(r"\\documentclass\[(.*?)\]{(.*?)}", r"\\documentclass[\1,fontset=windows,UTF8]{\2}",main_file)
- main_file = re.sub(r"\\documentclass{(.*?)}", r"\\documentclass[fontset=windows,UTF8]{\1}",main_file)
- # find paper abstract
- pattern_opt1 = re.compile(r'\\begin\{abstract\}.*\n')
- pattern_opt2 = re.compile(r"\\abstract\{(.*?)\}", flags=re.DOTALL)
- match_opt1 = pattern_opt1.search(main_file)
- match_opt2 = pattern_opt2.search(main_file)
- if (match_opt1 is None) and (match_opt2 is None):
- # "Cannot find paper abstract section!"
- main_file = insert_abstract(main_file)
- match_opt1 = pattern_opt1.search(main_file)
- match_opt2 = pattern_opt2.search(main_file)
- assert (match_opt1 is not None) or (match_opt2 is not None), "Cannot find paper abstract section!"
- return main_file
- insert_missing_abs_str = r"""
- \begin{abstract}
- The GPT-Academic program cannot find abstract section in this paper.
- \end{abstract}
- """
- def insert_abstract(tex_content):
- if "\\maketitle" in tex_content:
- # find the position of "\maketitle"
- find_index = tex_content.index("\\maketitle")
- # find the nearest ending line
- end_line_index = tex_content.find("\n", find_index)
- # insert "abs_str" on the next line
- modified_tex = tex_content[:end_line_index+1] + '\n\n' + insert_missing_abs_str + '\n\n' + tex_content[end_line_index+1:]
- return modified_tex
- elif r"\begin{document}" in tex_content:
- # find the position of "\maketitle"
- find_index = tex_content.index(r"\begin{document}")
- # find the nearest ending line
- end_line_index = tex_content.find("\n", find_index)
- # insert "abs_str" on the next line
- modified_tex = tex_content[:end_line_index+1] + '\n\n' + insert_missing_abs_str + '\n\n' + tex_content[end_line_index+1:]
- return modified_tex
- else:
- return tex_content
- """
- =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
- Post process
- =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
- """
- def mod_inbraket(match):
- """
- 为啥chatgpt会把cite里面的逗号换成中文逗号呀
- """
- # get the matched string
- cmd = match.group(1)
- str_to_modify = match.group(2)
- # modify the matched string
- str_to_modify = str_to_modify.replace(':', ':') # 前面是中文冒号,后面是英文冒号
- str_to_modify = str_to_modify.replace(',', ',') # 前面是中文逗号,后面是英文逗号
- # str_to_modify = 'BOOM'
- return "\\" + cmd + "{" + str_to_modify + "}"
- def fix_content(final_tex, node_string):
- """
- Fix common GPT errors to increase success rate
- """
- final_tex = re.sub(r"(?<!\\)%", "\\%", final_tex)
- final_tex = re.sub(r"\\([a-z]{2,10})\ \{", r"\\\1{", string=final_tex)
- final_tex = re.sub(r"\\\ ([a-z]{2,10})\{", r"\\\1{", string=final_tex)
- final_tex = re.sub(r"\\([a-z]{2,10})\{([^\}]*?)\}", mod_inbraket, string=final_tex)
- if "Traceback" in final_tex and "[Local Message]" in final_tex:
- final_tex = node_string # 出问题了,还原原文
- if node_string.count('\\begin') != final_tex.count('\\begin'):
- final_tex = node_string # 出问题了,还原原文
- if node_string.count('\_') > 0 and node_string.count('\_') > final_tex.count('\_'):
- # walk and replace any _ without \
- final_tex = re.sub(r"(?<!\\)_", "\\_", final_tex)
- def compute_brace_level(string):
- # this function count the number of { and }
- brace_level = 0
- for c in string:
- if c == "{": brace_level += 1
- elif c == "}": brace_level -= 1
- return brace_level
- def join_most(tex_t, tex_o):
- # this function join translated string and original string when something goes wrong
- p_t = 0
- p_o = 0
- def find_next(string, chars, begin):
- p = begin
- while p < len(string):
- if string[p] in chars: return p, string[p]
- p += 1
- return None, None
- while True:
- res1, char = find_next(tex_o, ['{','}'], p_o)
- if res1 is None: break
- res2, char = find_next(tex_t, [char], p_t)
- if res2 is None: break
- p_o = res1 + 1
- p_t = res2 + 1
- return tex_t[:p_t] + tex_o[p_o:]
- if compute_brace_level(final_tex) != compute_brace_level(node_string):
- # 出问题了,还原部分原文,保证括号正确
- final_tex = join_most(final_tex, node_string)
- return final_tex
-
- def compile_latex_with_timeout(command, cwd, timeout=60):
- import subprocess
- process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, cwd=cwd)
- try:
- stdout, stderr = process.communicate(timeout=timeout)
- except subprocess.TimeoutExpired:
- process.kill()
- stdout, stderr = process.communicate()
- print("Process timed out!")
- return False
- return True
- def run_in_subprocess_wrapper_func(func, args, kwargs, return_dict, exception_dict):
- import sys
- try:
- result = func(*args, **kwargs)
- return_dict['result'] = result
- except Exception as e:
- exc_info = sys.exc_info()
- exception_dict['exception'] = exc_info
- def run_in_subprocess(func):
- import multiprocessing
- def wrapper(*args, **kwargs):
- return_dict = multiprocessing.Manager().dict()
- exception_dict = multiprocessing.Manager().dict()
- process = multiprocessing.Process(target=run_in_subprocess_wrapper_func,
- args=(func, args, kwargs, return_dict, exception_dict))
- process.start()
- process.join()
- process.close()
- if 'exception' in exception_dict:
- # ooops, the subprocess ran into an exception
- exc_info = exception_dict['exception']
- raise exc_info[1].with_traceback(exc_info[2])
- if 'result' in return_dict.keys():
- # If the subprocess ran successfully, return the result
- return return_dict['result']
- return wrapper
- def _merge_pdfs(pdf1_path, pdf2_path, output_path):
- import PyPDF2 # PyPDF2这个库有严重的内存泄露问题,把它放到子进程中运行,从而方便内存的释放
- Percent = 0.95
- # raise RuntimeError('PyPDF2 has a serious memory leak problem, please use other tools to merge PDF files.')
- # Open the first PDF file
- with open(pdf1_path, 'rb') as pdf1_file:
- pdf1_reader = PyPDF2.PdfFileReader(pdf1_file)
- # Open the second PDF file
- with open(pdf2_path, 'rb') as pdf2_file:
- pdf2_reader = PyPDF2.PdfFileReader(pdf2_file)
- # Create a new PDF file to store the merged pages
- output_writer = PyPDF2.PdfFileWriter()
- # Determine the number of pages in each PDF file
- num_pages = max(pdf1_reader.numPages, pdf2_reader.numPages)
- # Merge the pages from the two PDF files
- for page_num in range(num_pages):
- # Add the page from the first PDF file
- if page_num < pdf1_reader.numPages:
- page1 = pdf1_reader.getPage(page_num)
- else:
- page1 = PyPDF2.PageObject.createBlankPage(pdf1_reader)
- # Add the page from the second PDF file
- if page_num < pdf2_reader.numPages:
- page2 = pdf2_reader.getPage(page_num)
- else:
- page2 = PyPDF2.PageObject.createBlankPage(pdf1_reader)
- # Create a new empty page with double width
- new_page = PyPDF2.PageObject.createBlankPage(
- width = int(int(page1.mediaBox.getWidth()) + int(page2.mediaBox.getWidth()) * Percent),
- height = max(page1.mediaBox.getHeight(), page2.mediaBox.getHeight())
- )
- new_page.mergeTranslatedPage(page1, 0, 0)
- new_page.mergeTranslatedPage(page2, int(int(page1.mediaBox.getWidth())-int(page2.mediaBox.getWidth())* (1-Percent)), 0)
- output_writer.addPage(new_page)
- # Save the merged PDF file
- with open(output_path, 'wb') as output_file:
- output_writer.write(output_file)
- merge_pdfs = run_in_subprocess(_merge_pdfs) # PyPDF2这个库有严重的内存泄露问题,把它放到子进程中运行,从而方便内存的释放
|