From da1982acb8b390926c7ee08308f23b33a73a2ccf Mon Sep 17 00:00:00 2001 From: Jake Poznanski Date: Mon, 30 Sep 2024 18:48:17 +0000 Subject: [PATCH] Refactoring prompts into their own new folder --- pdelfin/prompts/__init__.py | 1 + pdelfin/prompts/prompts.py | 24 ++++++++++++++++++++++++ pdelfin/silver_data/buildsilver.py | 16 ++-------------- pdelfin/train/dataprep.py | 14 ++------------ 4 files changed, 29 insertions(+), 26 deletions(-) create mode 100644 pdelfin/prompts/__init__.py create mode 100644 pdelfin/prompts/prompts.py diff --git a/pdelfin/prompts/__init__.py b/pdelfin/prompts/__init__.py new file mode 100644 index 0000000..2a4d945 --- /dev/null +++ b/pdelfin/prompts/__init__.py @@ -0,0 +1 @@ +from .prompts import * \ No newline at end of file diff --git a/pdelfin/prompts/prompts.py b/pdelfin/prompts/prompts.py new file mode 100644 index 0000000..30b83cd --- /dev/null +++ b/pdelfin/prompts/prompts.py @@ -0,0 +1,24 @@ +# This is the prompt we use for getting chat gpt 4o to convert documents into our silver training data +def build_openai_silver_data_prompt(base_text: str) -> str: + return ( + f"Below is the image of one page of a PDF document, as well as some raw textual content that was previously extracted for it. " + f"Just return the plain text representation of this document as if you were reading it naturally.\n" + f"Turn equations into a LaTeX representation, and tables into markdown format. Remove the headers and footers, but keep references and footnotes.\n" + f"Read any natural handwriting.\n" + f"This is likely one page out of several in the document, so be sure to preserve any sentences that come from the previous page, or continue onto the next page, exactly as they are.\n" + f"If there is no text at all that you think you should read, just output [NO TEXT].\n" + f"If the page has no English text on it at all, just output [NO ENGLISH TEXT].\n" + f"Do not hallucinate.\n" + f"RAW_TEXT_START\n{base_text}\nRAW_TEXT_END" + ) + + +# This is a base prompt that will be used for training and running the fine tuned model +# It's simplified from the prompt which was used to generate the silver data, and can change from dataset to dataset +def build_finetuning_prompt(base_text: str) -> str: + return ( + f"Below is the image of one page of a document, as well as some raw textual content that was previously extracted for it. " + f"Just return the plain text representation of this document as if you were reading it naturally.\n" + f"Do not hallucinate.\n" + f"RAW_TEXT_START\n{base_text}\nRAW_TEXT_END" + ) \ No newline at end of file diff --git a/pdelfin/silver_data/buildsilver.py b/pdelfin/silver_data/buildsilver.py index 1fc14a0..5d2282b 100644 --- a/pdelfin/silver_data/buildsilver.py +++ b/pdelfin/silver_data/buildsilver.py @@ -12,23 +12,11 @@ from typing import Generator from concurrent.futures import ThreadPoolExecutor, as_completed from urllib.parse import urlparse -# reuse mise pdf filtering base code +from pdelfin.prompts import build_openai_silver_data_prompt from pdelfin.filter import PdfFilter TARGET_IMAGE_DIM = 2048 -def _build_prompt(base_text: str) -> str: - return ( - f"Below is the image of one page of a PDF document, as well as some raw textual content that was previously extracted for it. " - f"Just return the plain text representation of this document as if you were reading it naturally.\n" - f"Turn equations into a LaTeX representation, and tables into markdown format. Remove the headers and footers, but keep references and footnotes.\n" - f"Read any natural handwriting.\n" - f"This is likely one page out of several in the document, so be sure to preserve any sentences that come from the previous page, or continue onto the next page, exactly as they are.\n" - f"If there is no text at all that you think you should read, just output [NO TEXT].\n" - f"If the page has no English text on it at all, just output [NO ENGLISH TEXT].\n" - f"Do not hallucinate.\n" - f"RAW_TEXT_START\n{base_text}\nRAW_TEXT_END" - ) pdf_filter = PdfFilter() @@ -78,7 +66,7 @@ def build_page_query(local_pdf_path: str, pretty_pdf_path: str, page: int) -> di { "role": "user", "content": [ - {"type": "text", "text": _build_prompt(base_text)}, + {"type": "text", "text": build_openai_silver_data_prompt(base_text)}, {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}} ], } diff --git a/pdelfin/train/dataprep.py b/pdelfin/train/dataprep.py index 53318b2..50f2a41 100644 --- a/pdelfin/train/dataprep.py +++ b/pdelfin/train/dataprep.py @@ -4,6 +4,7 @@ from PIL import Image import base64 import torch # Make sure to import torch as it's used in the DataCollator +from pdelfin.prompts import build_finetuning_prompt def filter_by_max_seq_len(example, processor, max_prompt_len: int=2000, max_response_len: int=2000): if len(processor.tokenizer.tokenize(example["input_prompt_text"])) > max_prompt_len: @@ -15,17 +16,6 @@ def filter_by_max_seq_len(example, processor, max_prompt_len: int=2000, max_resp return True -# This is a base prompt that will be used for training and running the fine tuned model -# It's simplified from the prompt which was used to generate the silver data, and can change from dataset to dataset -def _build_finetuning_prompt(base_text: str) -> str: - return ( - f"Below is the image of one page of a document, as well as some raw textual content that was previously extracted for it. " - f"Just return the plain text representation of this document as if you were reading it naturally.\n" - f"Do not hallucinate.\n" - f"RAW_TEXT_START\n{base_text}\nRAW_TEXT_END" - ) - - def prepare_data_for_qwen2_training(example, processor, add_batch_dim=False): # Prepare messages messages = [ @@ -36,7 +26,7 @@ def prepare_data_for_qwen2_training(example, processor, add_batch_dim=False): "type": "image", "image": example["input_prompt_image_base64"] # Placeholder }, - {"type": "text", "text": _build_finetuning_prompt(example["raw_page_text"])}, + {"type": "text", "text": build_finetuning_prompt(example["raw_page_text"])}, ], } ]