轻松解析本地PDF表格，基于LlamaIndex和UnstructuredIO打造RAG

from llama_index.readers.file.flat_reader import FlatReader
from llama_index.node_parser import UnstructuredElementNodeParser
import os
import pickle
from pathlib import Path

os.environ["OPENAI_API_KEY"] = "<your openai api key>"

# 读取数据
reader = FlatReader()
data = reader.load_data(Path('./The_Worlds_Billionaires.html'))

# 初始化 NodeParser
node_parser = UnstructuredElementNodeParser()

# 如果稍后想重用它
if not os.path.exists("qr_2023_nodes.pkl"):
    raw_nodes = node_parser.get_nodes_from_documents(data)
    pickle.dump(raw_nodes, open("the_world_billionaires_raw_nodes.pkl", "wb"))

# 基础节点和节点映射
base_nodes, node_mappings = node_parser.get_base_nodes_and_mappings(
    raw_nodes
)

from llama_index.retrievers import RecursiveRetriever
from llama_index.query_engine import RetrieverQueryEngine
from llama_index import VectorStoreIndex

vector_index = VectorStoreIndex(base_nodes_qr_2023)
vector_retriever = vector_index.as_retriever(similarity_top_k=3)
vector_query_engine = vector_index.as_query_engine(similarity_top_k=3)

recursive_retriever = RecursiveRetriever(
    "vector",
    retriever_dict={"vector": vector_retriever},
    node_dict=node_mappings_qr_2023,
)

query_engine = RetrieverQueryEngine.from_args(recursive_retriever)

query_engine.query("Who is the richest billionaire in 2020?")

from unstructured.partition.pdf import partition_pdf
from unstructured.staging.base import elements_to_json
import json

file_path = 'The_Worlds_Billionaires.pdf'
raw_pdf_elements = partition_pdf(
    filename=file_path,
    extract_images_in_pdf=False,
    infer_table_structure=True,
    chunking_strategy='by_title',
    max_characters=4000,
    new_after_n_chars=3800,
    combine_text_under_n_chars=2000,
    strategy = "hi_res"
)

# 将结果存储在 json 中
elements_to_json(raw_pdf_elements, filename=f"./The_Worlds_Billionaires_Converted.json")

no_tables = 0
def process_json_file(input_filename):
    # 读取 JSON 文件
    with open(f'./{input_filename}.json', 'r') as file:
        data = json.load(file)

    # 遍历 JSON 数据并提取所需的表格元素
    extracted_elements = []
    for entry in data:
        if entry['type'] == 'CompositeElement':
            extracted_elements.append(entry['text'])
        if entry["type"] == "Table":
            no_tables += 1
            extracted_elements.append(entry["metadata"]["text_as_html"])

    # 将提取的元素写入输出文件
    with open(f"{input_filename}.txt", 'w') as output_file:
        for element in extracted_elements:
            output_file.write(element + "\n\n")  # 添加两个换行符以分隔

process_json_file(f"The_Worlds_Billionaires_Converted")  # ## with new_file_name 是上面的 JSON 文件

print(f"Number of tables: {no_tables}")

### 加载数据

```python
# documents = SimpleDirectoryReader("./<folder_name>", 
#     input_files=['./<new_file_name.txt>']).load_data()