Intoval/privateChatGPT

1from types import SimpleNamespace2import pdfplumber3import logging4from llama_index import Document5 6def prepare_table_config(crop_page):7    """Prepare table查找边界, 要求page为原始page8    9    From https://github.com/jsvine/pdfplumber/issues/24210    """11    page = crop_page.root_page # root/parent12    cs = page.curves + page.edges13    def curves_to_edges():14        """See https://github.com/jsvine/pdfplumber/issues/127"""15        edges = []16        for c in cs:17            edges += pdfplumber.utils.rect_to_edges(c)18        return edges19    edges = curves_to_edges()20    return {21        "vertical_strategy": "explicit",22        "horizontal_strategy": "explicit",23        "explicit_vertical_lines": edges,24        "explicit_horizontal_lines": edges,25        "intersection_y_tolerance": 10,26    }27 28def get_text_outside_table(crop_page):29    ts = prepare_table_config(crop_page)30    if len(ts["explicit_vertical_lines"]) == 0 or len(ts["explicit_horizontal_lines"]) == 0:31        return crop_page32 33    ### Get the bounding boxes of the tables on the page.34    bboxes = [table.bbox for table in crop_page.root_page.find_tables(table_settings=ts)]35    def not_within_bboxes(obj):36        """Check if the object is in any of the table's bbox."""37        def obj_in_bbox(_bbox):38            """See https://github.com/jsvine/pdfplumber/blob/stable/pdfplumber/table.py#L404"""39            v_mid = (obj["top"] + obj["bottom"]) / 240            h_mid = (obj["x0"] + obj["x1"]) / 241            x0, top, x1, bottom = _bbox42            return (h_mid >= x0) and (h_mid < x1) and (v_mid >= top) and (v_mid < bottom)43        return not any(obj_in_bbox(__bbox) for __bbox in bboxes)44 45    return crop_page.filter(not_within_bboxes)46# 请使用 LaTeX 表达公式，行内公式以 $ 包裹，行间公式以 $$ 包裹47 48extract_words = lambda page: page.extract_words(keep_blank_chars=True, y_tolerance=0, x_tolerance=1, extra_attrs=["fontname", "size", "object_type"])49# dict_keys(['text', 'x0', 'x1', 'top', 'doctop', 'bottom', 'upright', 'direction', 'fontname', 'size'])50 51def get_title_with_cropped_page(first_page):52    title = [] # 处理标题53    x0,top,x1,bottom = first_page.bbox # 获取页面边框54 55    for word in extract_words(first_page):56        word = SimpleNamespace(**word)57 58        if word.size >= 14:59            title.append(word.text)60            title_bottom = word.bottom61        elif word.text == "Abstract": # 获取页面abstract62            top = word.top63    64    user_info = [i["text"] for i in extract_words(first_page.within_bbox((x0,title_bottom,x1,top)))]65    # 裁剪掉上半部分, within_bbox: full_included; crop: partial_included66    return title, user_info, first_page.within_bbox((x0,top,x1,bottom))67 68def get_column_cropped_pages(pages, two_column=True):69    new_pages = []70    for page in pages:71        if two_column:72            left = page.within_bbox((0, 0, page.width/2, page.height),relative=True)73            right = page.within_bbox((page.width/2, 0, page.width, page.height), relative=True)74            new_pages.append(left)75            new_pages.append(right)76        else:77            new_pages.append(page)78    79    return new_pages80 81def parse_pdf(filename, two_column = True):82    level = logging.getLogger().level83    if level == logging.getLevelName("DEBUG"):84        logging.getLogger().setLevel("INFO")85 86    with pdfplumber.open(filename) as pdf:87        title, user_info, first_page = get_title_with_cropped_page(pdf.pages[0])88        new_pages = get_column_cropped_pages([first_page] + pdf.pages[1:], two_column)89 90        chapters = []91        # tuple (chapter_name, [pageid] (start,stop), chapter_text)92        create_chapter = lambda page_start,name_top,name_bottom: SimpleNamespace(93            name=[],94            name_top=name_top,95            name_bottom=name_bottom,96            record_chapter_name = True,97            98            page_start=page_start,99            page_stop=None,100 101            text=[],102        )103        cur_chapter = None104 105        # 按页遍历PDF文档106        for idx, page in enumerate(new_pages):107            page = get_text_outside_table(page)108 109            # 按行遍历页面文本110            for word in extract_words(page):111                word = SimpleNamespace(**word)112 113                # 检查行文本是否以12号字体打印，如果是，则将其作为新章节开始114                if word.size >= 11: # 出现chapter name115                    if cur_chapter is None:116                        cur_chapter = create_chapter(page.page_number, word.top, word.bottom)117                    elif not cur_chapter.record_chapter_name or (cur_chapter.name_bottom != cur_chapter.name_bottom and cur_chapter.name_top != cur_chapter.name_top): 118                        # 不再继续写chapter name119                        cur_chapter.page_stop = page.page_number # stop id120                        chapters.append(cur_chapter)121                        # 重置当前chapter信息122                        cur_chapter = create_chapter(page.page_number, word.top, word.bottom)123 124                    # print(word.size, word.top, word.bottom, word.text)125                    cur_chapter.name.append(word.text)126                else:127                    cur_chapter.record_chapter_name = False # chapter name 结束128                    cur_chapter.text.append(word.text)129        else:130            # 处理最后一个章节131            cur_chapter.page_stop = page.page_number # stop id132            chapters.append(cur_chapter)133 134        for i in chapters:135            logging.info(f"section: {i.name} pages:{i.page_start, i.page_stop} word-count:{len(i.text)}")136            logging.debug(" ".join(i.text))137 138    title = " ".join(title)139    user_info = " ".join(user_info)140    text = f"Article Title: {title}, Information:{user_info}\n"141    for idx, chapter in enumerate(chapters):142        chapter.name = " ".join(chapter.name)143        text += f"The {idx}th Chapter {chapter.name}: " + " ".join(chapter.text) + "\n"144 145    logging.getLogger().setLevel(level)146    return Document(text=text, extra_info={"title": title})147 148BASE_POINTS = """1491. Who are the authors?1502. What is the process of the proposed method?1513. What is the performance of the proposed method? Please note down its performance metrics.1524. What are the baseline models and their performances? Please note down these baseline methods.1535. What dataset did this paper use?154"""155 156READING_PROMPT = """157You are a researcher helper bot. You can help the user with research paper reading and summarizing. \n158Now I am going to send you a paper. You need to read it and summarize it for me part by part. \n159When you are reading, You need to focus on these key points:{}160"""161 162READING_PROMT_V2 = """163You are a researcher helper bot. You can help the user with research paper reading and summarizing. \n164Now I am going to send you a paper. You need to read it and summarize it for me part by part. \n165When you are reading, You need to focus on these key points:{},166 167And You need to generate a brief but informative title for this part.168Your return format:169- title: '...'170- summary: '...'171"""172 173SUMMARY_PROMPT = "You are a researcher helper bot. Now you need to read the summaries of a research paper."174 175 176if __name__ == '__main__':177    # Test code178    z = parse_pdf("./build/test.pdf")179    print(z["user_info"])180    print(z["title"])