import asyncio import json import os import random import re from concurrent.futures import ThreadPoolExecutor from typing import Any, Dict, List, Optional import pandas as pd from datasets import IterableDataset, load_dataset from huggingface_hub import AsyncInferenceClient from tqdm import trange from tqdm.asyncio import tqdm_asyncio from transformers import AutoTokenizer from examples.question_answer_pairs.phase_1.base_prompts import (BASE_PROMPT, BASE_USER_CONTENT, PROMPTS) from llm_swarm import LLMSwarm, LLMSwarmConfig CHECKPOINT_FILE = 'checkpoint.json' DATA_PATH = '/fsx/andi/pdfa_data/' TAR_FILE_PATTERN = 'pdfa-eng-train-{:06d}.tar' NUM_TAR_FILES = 1800 # Total number of tar files MAX_PAGES_PER_PDF = 4 STEP_SIZE = 10 model_id = "microsoft/Phi-3-small-8k-instruct" def create_llm_prompt(prompt, text): system_content = BASE_PROMPT.format( role_description=prompt["role_description"], examples=prompt["examples"] ) return [ {"role": "system", "content": system_content}, {"role": "user", "content": BASE_USER_CONTENT.format(text=text)} ] def extract_text_per_page_from_sample(sample: Dict[str, Any]) -> List[str]: """ Extracts text from each page of a given sample and returns it as a list of strings. Args: sample (Dict[str, Any]): The sample containing page data in JSON format. Returns: List[str]: A list of strings, where each string represents the text of a page. """ texts = [] for page in sample['json']['pages']: pages_text = ' \n '.join(page['lines']['text']) texts.append(pages_text) return texts def extract_chunks(pages: List[Any], max_tokens_per_group: int, max_pages_per_group: int, n_overlap: int) -> List[str]: """ Splits a list of pages into chunks with a specified maximum number of tokens per chunk, a maximum number of pages per chunk, and overlap between chunks. Args: pages (List[Any]): The list of pages to be chunked. max_tokens_per_group (int): The maximum number of tokens allowed per chunk. max_pages_per_group (int): The maximum number of pages allowed per chunk. n_overlap (int): The number of overlapping pages between consecutive chunks. Returns: List[str]: A list of chunked text, each chunk containing text from multiple pages. """ chunks = [] current_chunk = [] current_chunk_tokens = 0 current_chunk_pages = 0 page_token_counts = [len(tokenizer.encode(page, add_special_tokens=False)) for page in pages] for i, page in enumerate(pages): page_tokens = page_token_counts[i] if page_tokens > max_tokens_per_group: print(f"Skipping document where page nr {i} has {page_tokens} tokens.") return [] if (current_chunk_tokens + page_tokens > max_tokens_per_group) or (current_chunk_pages + 1 > max_pages_per_group): if current_chunk: chunks.append('\nNEW PAGE\n'.join(current_chunk)) current_chunk = current_chunk[-n_overlap:] if n_overlap > 0 else [] current_chunk_tokens = sum(page_token_counts[max(0, i - n_overlap):i]) current_chunk_pages = len(current_chunk) current_chunk.append(page) current_chunk_tokens += page_tokens current_chunk_pages += 1 if current_chunk: chunks.append('\nNEW PAGE\n'.join(current_chunk)) return chunks def create_tasks(dataset: IterableDataset, prompt_id: Optional[int] = None, n_overlap: int = 2) -> List[Dict[str, Any]]: """ Processes a dataset to generate question and answer pairs for each sample. Args: dataset (IterableDataset): The dataset containing samples. prompt_id (Optional[int]): The ID of the prompt template to use for generating questions. If set to None, prompt_id is random. n_overlap (int): The number of overlapping pages between consecutive chunks. num_samples (int): The number of samples to process. Returns: List[Dict[str, Any]]: A list of dictionaries containing the sample key, page count, generated Q/A pairs, and prompt ID. """ if prompt_id is not None: selected_id_prompt = prompt_id tasks = [] for index, sample in dataset.iterrows(): text_per_page = extract_text_per_page_from_sample(sample) if len(text_per_page) > MAX_PAGES_PER_PDF: continue page_chunks = extract_chunks(text_per_page, max_tokens_per_group=5000, max_pages_per_group=5, n_overlap=n_overlap) for chunk in page_chunks: if prompt_id is None: selected_id_prompt = random.randint(0, 4) prompt = PROMPTS[selected_id_prompt] messages = create_llm_prompt(prompt, chunk) prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True ) tasks_dict = { "__key__": sample['__key__'], "Page count": len(text_per_page), "messages": prompt, "Prompt ID": selected_id_prompt } tasks.append(tasks_dict) return tasks # Function to extract Q&A pairs from a string def extract_qa_pairs(text): qa_pattern = re.compile(r'(Q\d+:\s*.*?)(A\d+:\s*.*?)(?=(Q\d+:)|$)', re.DOTALL) matches = qa_pattern.findall(text) qa_pairs = [(q.strip(), a.strip()) for match in matches for q, a in [match[:2]]] return qa_pairs def process_outputs_to_df(df): all_data = [] for index, row in df.iterrows(): task = row['Task'] completion = row['Completion'] sample_key = task['__key__'] page_count = task['Page count'] prompt_id = task['Prompt ID'] qa_pairs = extract_qa_pairs(completion) if len(qa_pairs) == 0: print('No Q&A pairs found for sample:', sample_key) for question, answer in qa_pairs: all_data.append({ '__key__': sample_key, 'Page count': page_count, 'Prompt ID': prompt_id, 'question': question, 'answer': answer }) qa_df = pd.DataFrame(all_data) return qa_df def save_checkpoint(tar_index, total_examples): checkpoint_data = { 'tar_index': tar_index, 'total_examples': total_examples } with open(CHECKPOINT_FILE, 'w') as f: json.dump(checkpoint_data, f) def load_checkpoint(): if os.path.exists(CHECKPOINT_FILE): with open(CHECKPOINT_FILE, 'r') as f: return json.load(f) return {'tar_index': 0, 'total_examples': 0} tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) def launch(): with LLMSwarm( LLMSwarmConfig( instances=8, inference_engine="vllm", gpus=1, model=model_id, slurm_template_path="templates/vllm_h100.template.slurm", load_balancer_template_path="templates/nginx.template.conf", trust_remote_code=True, per_instance_max_parallel_requests=200, ) ) as llm_swarm: semaphore = asyncio.Semaphore(llm_swarm.suggested_max_parallel_requests) client = AsyncInferenceClient(model=llm_swarm.endpoint) async def process_text(prompt): async with semaphore: response = await client.post( json={ "prompt": prompt, "max_tokens": 2000, } ) res = json.loads(response.decode("utf-8"))["text"][0][len(prompt):] return res def load_and_process_dataset(tar_file): try: print(f"Loading dataset from: {tar_file}") dataset = load_dataset('webdataset', split='train', data_files=tar_file).to_pandas() tasks = create_tasks(dataset, prompt_id=None, n_overlap=1) return tasks except Exception as e: print(f"Error loading dataset from: {tar_file}") print(e) return [] def get_future_tasks(tar_index, executor): futures = [] for inner_idx in range(STEP_SIZE): tar_file = os.path.join(DATA_PATH, TAR_FILE_PATTERN.format(tar_index + inner_idx)) futures.append(executor.submit(load_and_process_dataset, tar_file)) return futures async def process_dataset(tar_index, total_examples): next_future_tasks = get_future_tasks(tar_index, ThreadPoolExecutor(max_workers=STEP_SIZE)) for idx in trange(tar_index, NUM_TAR_FILES + STEP_SIZE, STEP_SIZE, desc="Creating Dataset"): print(f"Processing tar file {idx}") tasks = [] future_tasks = next_future_tasks results = [f.result() for f in future_tasks] for result in results: tasks.extend(result) # Once you created the tasks for this batch, load the next batch in parallel # Otherwise, the tasks for this batch compete with the tasks from next batch for resources next_future_tasks = get_future_tasks(idx + STEP_SIZE, ThreadPoolExecutor(max_workers=1)) # Only one thread to avoid cpu clogging results = await tqdm_asyncio.gather(*(process_text(task['messages']) for task in tasks)) df = pd.DataFrame({"Task": tasks, "Completion": results}) df_new = process_outputs_to_df(df) # Save the batch to HDF5 df_new.to_hdf(f'synthetic_dataset_batch_{idx}.h5', key='df', mode='w') unique_keys = df_new['__key__'].nunique() total_examples += unique_keys save_checkpoint(idx, total_examples) async def main(): checkpoint = load_checkpoint() tar_index = checkpoint['tar_index'] if tar_index != 0: tar_index += STEP_SIZE print(f"Resuming from tar file {tar_index}") total_examples = checkpoint['total_examples'] processor = asyncio.create_task(process_dataset(tar_index, total_examples)) await processor print("All batches processed.") asyncio.run(main()) launch()