modules/pdf_func.py

from types import SimpleNamespace
import pdfplumber
import logging
from llama_index import Document

def prepare_table_config(crop_page):
    """Prepare table查找边界, 要求page为原始page
    
    From https://github.com/jsvine/pdfplumber/issues/242
    """
    page = crop_page.root_page # root/parent
    cs = page.curves + page.edges
    def curves_to_edges():
        """See https://github.com/jsvine/pdfplumber/issues/127"""
        edges = []
        for c in cs:
            edges += pdfplumber.utils.rect_to_edges(c)
        return edges
    edges = curves_to_edges()
    return {
        "vertical_strategy": "explicit",
        "horizontal_strategy": "explicit",
        "explicit_vertical_lines": edges,
        "explicit_horizontal_lines": edges,
        "intersection_y_tolerance": 10,
    }

def get_text_outside_table(crop_page):
    ts = prepare_table_config(crop_page)
    if len(ts["explicit_vertical_lines"]) == 0 or len(ts["explicit_horizontal_lines"]) == 0:
        return crop_page

    ### Get the bounding boxes of the tables on the page.
    bboxes = [table.bbox for table in crop_page.root_page.find_tables(table_settings=ts)]
    def not_within_bboxes(obj):
        """Check if the object is in any of the table's bbox."""
        def obj_in_bbox(_bbox):
            """See https://github.com/jsvine/pdfplumber/blob/stable/pdfplumber/table.py#L404"""
            v_mid = (obj["top"] + obj["bottom"]) / 2
            h_mid = (obj["x0"] + obj["x1"]) / 2
            x0, top, x1, bottom = _bbox
            return (h_mid >= x0) and (h_mid < x1) and (v_mid >= top) and (v_mid < bottom)
        return not any(obj_in_bbox(__bbox) for __bbox in bboxes)

    return crop_page.filter(not_within_bboxes)
# 请使用 LaTeX 表达公式，行内公式以 $ 包裹，行间公式以 $$ 包裹

extract_words = lambda page: page.extract_words(keep_blank_chars=True, y_tolerance=0, x_tolerance=1, extra_attrs=["fontname", "size", "object_type"])
# dict_keys(['text', 'x0', 'x1', 'top', 'doctop', 'bottom', 'upright', 'direction', 'fontname', 'size'])

def get_title_with_cropped_page(first_page):
    title = [] # 处理标题
    x0,top,x1,bottom = first_page.bbox # 获取页面边框

    for word in extract_words(first_page):
        word = SimpleNamespace(**word)

        if word.size >= 14:
            title.append(word.text)
            title_bottom = word.bottom
        elif word.text == "Abstract": # 获取页面abstract
            top = word.top
    
    user_info = [i["text"] for i in extract_words(first_page.within_bbox((x0,title_bottom,x1,top)))]
    # 裁剪掉上半部分, within_bbox: full_included; crop: partial_included
    return title, user_info, first_page.within_bbox((x0,top,x1,bottom))

def get_column_cropped_pages(pages, two_column=True):
    new_pages = []
    for page in pages:
        if two_column:
            left = page.within_bbox((0, 0, page.width/2, page.height),relative=True)
            right = page.within_bbox((page.width/2, 0, page.width, page.height), relative=True)
            new_pages.append(left)
            new_pages.append(right)
        else:
            new_pages.append(page)
    
    return new_pages

def parse_pdf(filename, two_column = True):
    level = logging.getLogger().level
    if level == logging.getLevelName("DEBUG"):
        logging.getLogger().setLevel("INFO")

    with pdfplumber.open(filename) as pdf:
        title, user_info, first_page = get_title_with_cropped_page(pdf.pages[0])
        new_pages = get_column_cropped_pages([first_page] + pdf.pages[1:], two_column)

        chapters = []
        # tuple (chapter_name, [pageid] (start,stop), chapter_text)
        create_chapter = lambda page_start,name_top,name_bottom: SimpleNamespace(
            name=[],
            name_top=name_top,
            name_bottom=name_bottom,
            record_chapter_name = True,
            
            page_start=page_start,
            page_stop=None,

            text=[],
        )
        cur_chapter = None

        # 按页遍历PDF文档
        for idx, page in enumerate(new_pages):
            page = get_text_outside_table(page)

            # 按行遍历页面文本
            for word in extract_words(page):
                word = SimpleNamespace(**word)

                # 检查行文本是否以12号字体打印，如果是，则将其作为新章节开始
                if word.size >= 11: # 出现chapter name
                    if cur_chapter is None:
                        cur_chapter = create_chapter(page.page_number, word.top, word.bottom)
                    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): 
                        # 不再继续写chapter name
                        cur_chapter.page_stop = page.page_number # stop id
                        chapters.append(cur_chapter)
                        # 重置当前chapter信息
                        cur_chapter = create_chapter(page.page_number, word.top, word.bottom)

                    # print(word.size, word.top, word.bottom, word.text)
                    cur_chapter.name.append(word.text)
                else:
                    cur_chapter.record_chapter_name = False # chapter name 结束
                    cur_chapter.text.append(word.text)
        else:
            # 处理最后一个章节
            cur_chapter.page_stop = page.page_number # stop id
            chapters.append(cur_chapter)

        for i in chapters:
            logging.info(f"section: {i.name} pages:{i.page_start, i.page_stop} word-count:{len(i.text)}")
            logging.debug(" ".join(i.text))

    title = " ".join(title)
    user_info = " ".join(user_info)
    text = f"Article Title: {title}, Information:{user_info}\n"
    for idx, chapter in enumerate(chapters):
        chapter.name = " ".join(chapter.name)
        text += f"The {idx}th Chapter {chapter.name}: " + " ".join(chapter.text) + "\n"

    logging.getLogger().setLevel(level)
    return Document(text=text, extra_info={"title": title})

BASE_POINTS = """
1. Who are the authors?
2. What is the process of the proposed method?
3. What is the performance of the proposed method? Please note down its performance metrics.
4. What are the baseline models and their performances? Please note down these baseline methods.
5. What dataset did this paper use?
"""

READING_PROMPT = """
You are a researcher helper bot. You can help the user with research paper reading and summarizing. \n
Now I am going to send you a paper. You need to read it and summarize it for me part by part. \n
When you are reading, You need to focus on these key points:{}
"""

READING_PROMT_V2 = """
You are a researcher helper bot. You can help the user with research paper reading and summarizing. \n
Now I am going to send you a paper. You need to read it and summarize it for me part by part. \n
When you are reading, You need to focus on these key points:{},

And You need to generate a brief but informative title for this part.
Your return format:
- title: '...'
- summary: '...'
"""

SUMMARY_PROMPT = "You are a researcher helper bot. Now you need to read the summaries of a research paper."


if __name__ == '__main__':
    # Test code
    z = parse_pdf("./build/test.pdf")
    print(z["user_info"])
    print(z["title"])