zh_CN/Step-1-PDF预处理.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4f67a6a6",
   "metadata": {},
   "source": [
    "## Notebook 1: PDF 预处理"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30b6d43441645f5b",
   "metadata": {},
   "source": [
    "在这一系列中，我们将使用开源模型将PDF转换为播客。\n",
    "\n",
    "获取播客的第一步是找到一个脚本，目前我们的步骤如下：\n",
    "- 使用任意主题的PDF\n",
    "- 使用 LLM 模型，将其处理成文本文件。英文用的是Qwen2.5-1.5B，中文使用Qwen2.5-1.5B效果不太好。\n",
    "- 在下一个 Notebook 中将其改写为播客稿件。\n",
    "\n",
    "在这个 Notebook 中，我们将上传一个PDF，使用 `PyPDF2` 将其保存为 `.txt` 文件，后续我们会用轻量模型处理文本。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "582680da8eac6c30",
   "metadata": {},
   "source": [
    "务必记得安装此项目依赖，不然跑不起来～🙂‍↔️"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f4fc7aef-3505-482e-a998-790b8b9d48e4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-31T16:04:58.318414Z",
     "start_time": "2024-10-31T16:04:58.316740Z"
    }
   },
   "outputs": [],
   "source": [
    "#!pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b23d509",
   "metadata": {},
   "source": [
    "这里设置需要处理的PDF文件路径。\n",
    "\n",
    "另外，如果你想要发挥GPU极致性能，求你更换大参数的模型，虽然轻量级模型对于此任务也能胜任。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "60d0061b-8b8c-4353-850f-f19466a0ae2d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-31T16:04:58.339613Z",
     "start_time": "2024-10-31T16:04:58.337923Z"
    }
   },
   "outputs": [],
   "source": [
    "pdf_path = '../resources/2402.13116v4.pdf'\n",
    "DEFAULT_MODEL = \"llama3\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "21029232-ac5f-42ca-b26b-baad5b2f49b7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-31T16:04:58.351074Z",
     "start_time": "2024-10-31T16:04:58.349366Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import warnings\n",
    "from typing import Optional\n",
    "\n",
    "import PyPDF2\n",
    "from mlx_lm import load, generate\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "203c22eb",
   "metadata": {},
   "source": [
    "这里会检查文件有没有什么毛病～"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "153d9ece-37a4-4fff-a8e8-53f923a2b0a0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-31T16:04:58.366385Z",
     "start_time": "2024-10-31T16:04:58.364431Z"
    }
   },
   "outputs": [],
   "source": [
    "def validate_pdf(file_path: str) -> bool:\n",
    "    if not os.path.exists(file_path):\n",
    "        print(f\"Error: 此文件找不到啊🤔: {file_path}\")\n",
    "        return False\n",
    "    if not file_path.lower().endswith('.pdf'):\n",
    "        print(\"Error: 此文件不是 PDF 😵‍💫\")\n",
    "        return False\n",
    "    return True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a362ac3",
   "metadata": {},
   "source": [
    "这里就简单的读取 PDF 保存为`.txt`文件。默认的最大字符数是10万，如果你有长篇大论要处理，随时可以调整哦！不过，记得要考虑模型的长度限制～"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "b57c2d64-3d75-4aeb-b4ee-bd1661286b66",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-31T16:04:58.376960Z",
     "start_time": "2024-10-31T16:04:58.373410Z"
    }
   },
   "outputs": [],
   "source": [
    "def extract_text_from_pdf(file_path: str, max_chars: int = 100000) -> Optional[str]:\n",
    "    if not validate_pdf(file_path):\n",
    "        return None\n",
    "\n",
    "    try:\n",
    "        with open(file_path, 'rb') as file:\n",
    "            # 创建 PDF reader 对象\n",
    "            pdf_reader = PyPDF2.PdfReader(file)\n",
    "\n",
    "            # 获取总页数\n",
    "            num_pages = len(pdf_reader.pages)\n",
    "            print(f\"处理 {num_pages} 页的 PDF...\")\n",
    "\n",
    "            extracted_text = []\n",
    "            total_chars = 0\n",
    "\n",
    "            # 遍历所有页面\n",
    "            for page_num in range(num_pages):\n",
    "                # 提取当前页面的文本\n",
    "                page = pdf_reader.pages[page_num]\n",
    "                text = page.extract_text()\n",
    "\n",
    "                # 检查文本是否超过最大字符数\n",
    "                if total_chars + len(text) > max_chars:\n",
    "                    # 仅截取到最大字符数\n",
    "                    remaining_chars = max_chars - total_chars\n",
    "                    extracted_text.append(text[:remaining_chars])\n",
    "                    print(f\"在第 {page_num + 1} 页超过 {max_chars} 字符限制啦！！\")\n",
    "                    break\n",
    "\n",
    "                extracted_text.append(text)\n",
    "                total_chars += len(text)\n",
    "                print(f\"已处理 {page_num + 1}/{num_pages} 页\")\n",
    "\n",
    "            final_text = '\\n'.join(extracted_text)\n",
    "            print(f\"\\n提取完成！总字数：{len(final_text)}\")\n",
    "            return final_text\n",
    "\n",
    "    except PyPDF2.PdfReadError:\n",
    "        print(\"Error: 这 PDF 有毛病啊，请仔细看看～\")\n",
    "        return None\n",
    "    except Exception as e:\n",
    "        print(f\"完了，挂了: {str(e)}\")\n",
    "        return None\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e023397b",
   "metadata": {},
   "source": [
    "获取 PDF 元信息"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "0984bb1e-d52c-4cec-a131-67a48061fabc",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-31T16:04:58.386444Z",
     "start_time": "2024-10-31T16:04:58.384045Z"
    }
   },
   "outputs": [],
   "source": [
    "# 获取 PDF 元信息\n",
    "def get_pdf_metadata(file_path: str) -> Optional[dict]:\n",
    "    if not validate_pdf(file_path):\n",
    "        return None\n",
    "\n",
    "    try:\n",
    "        with open(file_path, 'rb') as file:\n",
    "            pdf_reader = PyPDF2.PdfReader(file)\n",
    "            metadata = {\n",
    "                'num_pages': len(pdf_reader.pages),\n",
    "                'metadata': pdf_reader.metadata\n",
    "            }\n",
    "            return metadata\n",
    "    except Exception as e:\n",
    "        print(f\"提取元信息失败: {str(e)}\")\n",
    "        return None"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6019affc",
   "metadata": {},
   "source": [
    "Finally, we can run our logic to extract the details from the file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63848943-79cc-4e21-8396-6eab5df493e0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-31T16:04:58.801603Z",
     "start_time": "2024-10-31T16:04:58.392756Z"
    }
   },
   "outputs": [],
   "source": [
    "# 提取元信息\n",
    "print(\"正在提取元信息...\")\n",
    "metadata = get_pdf_metadata(pdf_path)\n",
    "if metadata:\n",
    "    print(\"\\nPDF 元信息:\")\n",
    "    print(f\"页数: {metadata['num_pages']}\")\n",
    "    print(\"文档信息:\")\n",
    "    for key, value in metadata['metadata'].items():\n",
    "        print(f\"{key}: {value}\")\n",
    "\n",
    "# 提取文本\n",
    "print(\"\\n提取文本...\")\n",
    "extracted_text = extract_text_from_pdf(pdf_path)\n",
    "\n",
    "# 预览前500字符文本\n",
    "if extracted_text:\n",
    "    print(\"\\n预览提取后的文本 (前500字符):\")\n",
    "    print(\"-\" * 50)\n",
    "    print(extracted_text[:500])\n",
    "    print(\"-\" * 50)\n",
    "    print(f\"\\n提取的总字符数: {len(extracted_text)}\")\n",
    "\n",
    "# 将提取的文本保存到 txt 文件中\n",
    "if extracted_text:\n",
    "    output_file = './resources/extracted_text.txt'\n",
    "    with open(output_file, 'w', encoding='utf-8') as f:\n",
    "        f.write(extracted_text)\n",
    "    print(f\"\\n提取的文本已保存到 {output_file}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "946d1f59",
   "metadata": {},
   "source": [
    "### 预处理\n",
    "现在我们来说说为啥不用正则表达式的提取，而使用大型语言模型的：\n",
    "目前，我们已经从 PDF 中提取出一个文本文件。一般来说，由于字符、格式、Latex、表格等原因，PDF 提取出的内容可能会很凌乱。\n",
    "解决这个问题的一种方法是使用正则表达式，但我们也可以使用轻量模型帮助清洗文本。\n",
    "你可以试试修改 `SYS_PROMPT`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "7c0828a5-964d-475e-b5f5-40a04e287725",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-31T16:04:58.809790Z",
     "start_time": "2024-10-31T16:04:58.807967Z"
    }
   },
   "outputs": [],
   "source": [
    "SYS_PROMPT = \"\"\"您是一流的文本预处理器。以下是来自 PDF 的原始数据，请解析并以清晰可用的方式返回给播客作者。\n",
    "原始数据中包含换行符、Latex 数学公式以及一些可以完全删除的冗余内容。请去掉对播客作者记录无用的信息。\n",
    "不论主题如何，以上问题并不详尽，因此请聪明选择要删除的信息，发挥创造力。\n",
    "请注意，您的任务仅限于清理文本和必要时重写，在删除细节时要果断。当您进行文本处理时，请确保输出中文。\n",
    "如果遇到英文或其他语言，请翻译为中文后再输出。\n",
    "请勿添加 Markdown 格式化或特殊字符。\n",
    "一开始便直接开始响应处理后的文本，不做任何确认，谢谢！\n",
    "这里是文本：\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd393fae",
   "metadata": {},
   "source": [
    "为了避免模型一次处理整个文件，我们将分块处理文件。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "24e8a547-9d7c-4e2f-be9e-a3aea09cce76",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-31T16:04:58.818065Z",
     "start_time": "2024-10-31T16:04:58.815695Z"
    }
   },
   "outputs": [],
   "source": [
    "def create_word_bounded_chunks(text, target_chunk_size):\n",
    "    \"\"\"\n",
    "    这里会根据 target_chunk_size 分块\n",
    "    \"\"\"\n",
    "    words = text.split()\n",
    "    chunks = []\n",
    "    current_chunk = []\n",
    "    current_length = 0\n",
    "\n",
    "    for word in words:\n",
    "        word_length = len(word) + 1  # +1 for the space\n",
    "        if current_length + word_length > target_chunk_size and current_chunk:\n",
    "            # Join the current chunk and add it to chunks\n",
    "            chunks.append(' '.join(current_chunk))\n",
    "            current_chunk = [word]\n",
    "            current_length = word_length\n",
    "        else:\n",
    "            current_chunk.append(word)\n",
    "            current_length += word_length\n",
    "\n",
    "    # Add the last chunk if it exists\n",
    "    if current_chunk:\n",
    "        chunks.append(' '.join(current_chunk))\n",
    "\n",
    "    return chunks"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d74223f",
   "metadata": {},
   "source": [
    "让我们开始处理文本块吧！✨"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d3519667",
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "openapi_client = openai.Client(base_url=\"http://localhost:8080/v1\",api_key='NA')\n",
    "\n",
    "def process_chunk_with_webapi(text_chunk, chunk_num):\n",
    "    \"\"\"处理文本块，并返回模型处理好的文本\"\"\"\n",
    "    conversation = [\n",
    "        {\"role\": \"system\", \"content\": SYS_PROMPT},\n",
    "        {\"role\": \"user\", \"content\": text_chunk},\n",
    "    ]\n",
    "\n",
    "    response = openapi_client.chat.completions.create(\n",
    "        messages=conversation,\n",
    "        model=DEFAULT_MODEL\n",
    "    )\n",
    "    processed_text = response.choices[0].message.content\n",
    "\n",
    "    # Print chunk information for monitoring\n",
    "    #print(f\"\\n{'='*40} Chunk {chunk_num} {'='*40}\")\n",
    "    print(f\"INPUT TEXT:\\n{text_chunk[:500]}...\")  # Show first 500 chars of input\n",
    "    print(f\"\\nPROCESSED TEXT:\\n{processed_text[:500]}...\")  # Show first 500 chars of output\n",
    "    print(f\"{'=' * 90}\\n\")\n",
    "\n",
    "    return processed_text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "a0183c47-339d-4041-ae83-77fc34931075",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-31T16:05:01.192345Z",
     "start_time": "2024-10-31T16:05:01.187401Z"
    }
   },
   "outputs": [],
   "source": [
    "INPUT_FILE = \"./resources/extracted_text.txt\"  \n",
    "CHUNK_SIZE = 1000\n",
    "\n",
    "# 读取文件\n",
    "with open(INPUT_FILE, 'r', encoding='utf-8') as file:\n",
    "    text = file.read()\n",
    "\n",
    "chunks = create_word_bounded_chunks(extracted_text, CHUNK_SIZE)\n",
    "num_chunks = len(chunks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "bb36814f-9310-4734-bf54-e16a5032339e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-31T16:05:01.200630Z",
     "start_time": "2024-10-31T16:05:01.198749Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "101"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "num_chunks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "59d994601d77baa8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-31T16:05:01.311376Z",
     "start_time": "2024-10-31T16:05:01.309706Z"
    }
   },
   "outputs": [],
   "source": [
    "# 创建输出文件路径\n",
    "output_file = f\"./resources/clean_{os.path.basename(INPUT_FILE)}\"\n",
    "\n",
    "processed_text = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7917dfdd-b3af-44fc-a8c0-2760ace9363e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-31T16:12:11.049886Z",
     "start_time": "2024-10-31T16:05:01.390828Z"
    }
   },
   "outputs": [],
   "source": [
    "with open(output_file, 'w', encoding='utf-8') as out_file:\n",
    "    for chunk_num, chunk in enumerate(tqdm(chunks, desc=\"Processing chunks\")):\n",
    "        # 处理文本块\n",
    "        processed_chunk = process_chunk_with_webapi(chunk, chunk_num)\n",
    "        processed_text += processed_chunk + \"\\n\"\n",
    "\n",
    "        # 将处理后的文本写入输出文件\n",
    "        out_file.write(processed_chunk + \"\\n\")\n",
    "        out_file.flush()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31cffe8d",
   "metadata": {},
   "source": [
    "让我们看看最后处理的效果吧～🍻"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89ef51a7-f13f-49a4-8f73-9ac8ce75319d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-31T16:12:11.074075Z",
     "start_time": "2024-10-31T16:12:11.071563Z"
    }
   },
   "outputs": [],
   "source": [
    "print(f\"\\n处理完成！\")\n",
    "print(f\"输入文件: {INPUT_FILE}\")\n",
    "print(f\"输出文件: {output_file}\")\n",
    "print(f\"已处理总文本块: {num_chunks}\")\n",
    "\n",
    "# 预览处理后的文本的开头和结尾。\n",
    "print(\"\\n预览处理后的文本：\")\n",
    "print(\"\\n开头:\")\n",
    "print(processed_text[:1000])\n",
    "print(\"\\n...\\n\\n结尾:\")\n",
    "print(processed_text[-1000:])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d996ac5",
   "metadata": {},
   "source": [
    "### 下一个 Notebook: 转录员\n",
    "\n",
    "现在我们已经预处理好文本，在下一个 Notebook 中将其转换为讲稿"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "1b16ae0e-04cf-4eb9-a369-dee1728b89ce",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-31T16:12:11.092996Z",
     "start_time": "2024-10-31T16:12:11.091314Z"
    }
   },
   "outputs": [],
   "source": [
    "#fin"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}