From 302ed4dbddf46699dddba395fb71b0a8d5e112fe Mon Sep 17 00:00:00 2001 From: 51616 Date: Mon, 13 Jan 2025 14:21:41 +0000 Subject: [PATCH] include new data! --- .gitignore | 1 + README.md | 8 + configs/fw_qa_tiny.yaml | 43 +++ configs/pwc_hotpot_qa.yaml | 7 +- generate_fw_qa.py | 162 +++++++++ generate_qa_parallel.py | 493 ++++++++++++++++++++++++++++ hyperlora/configs.py | 30 +- hyperlora/data_utils.py | 52 +-- hyperlora/intx_sft.py | 33 +- install.sh | 14 - post_process_fw_qa.py | 216 ++++++++++++ post_process_parallel_qa.py | 206 ++++++++++++ process_automathtext_arxiv.py | 92 ++++++ process_automathtext_code_python.py | 93 ++++++ process_automathtext_web.py | 92 ++++++ process_codeparrot_qa.py | 90 +++++ process_fineweb.py | 37 +++ process_openphi_prog.py | 91 +++++ 18 files changed, 1710 insertions(+), 50 deletions(-) create mode 100644 configs/fw_qa_tiny.yaml create mode 100644 generate_fw_qa.py create mode 100644 generate_qa_parallel.py mode change 100644 => 100755 hyperlora/intx_sft.py delete mode 100755 install.sh create mode 100644 post_process_fw_qa.py create mode 100644 post_process_parallel_qa.py create mode 100644 process_automathtext_arxiv.py create mode 100644 process_automathtext_code_python.py create mode 100644 process_automathtext_web.py create mode 100644 process_codeparrot_qa.py create mode 100644 process_fineweb.py create mode 100644 process_openphi_prog.py diff --git a/.gitignore b/.gitignore index c7c2960..5ed6ddb 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ .cursorrules tmp/ +openai_batches/ models/ results/ datasets/ diff --git a/README.md b/README.md index c7c420d..a3a99c6 100644 --- a/README.md +++ b/README.md @@ -14,4 +14,12 @@ WANDB_MODE=disabled run python hyperlora/intx_sft.py configs/context_numbers_10. ### HyperLoRA w/ context_numbers_128 ```bash WANDB_MODE=disabled run python hyperlora/intx_sft.py configs/context_numbers_128.yaml --model_name_or_path=meta-llama/Llama-3.2-1B-Instruct --num_train_epochs=10 --per_device_train_batch_size=64 --per_device_eval_batch_size=8 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj +``` + +### Generate fineweb qa +```bash +# this might take several days... +python process_fineweb.py +python generate_fw_qa.py +python post_process_fw_qa.py ``` \ No newline at end of file diff --git a/configs/fw_qa_tiny.yaml b/configs/fw_qa_tiny.yaml new file mode 100644 index 0000000..f3d359d --- /dev/null +++ b/configs/fw_qa_tiny.yaml @@ -0,0 +1,43 @@ +output_dir: "" # just a placeholder +bf16: true +model_name_or_path: meta-llama/Llama-3.2-1B-Instruct +label_names: ["labels"] +# eval_on_start: True +# eval_strategy: "steps" +# eval_steps: 500 +# save_strategy: "no" +# # save_steps: 500 +# logging_strategy: "steps" +# logging_steps: 100 +# use_liger_kernel: true +# remove_unused_columns: false + +# needed to avoid OOM by compute the metrics batch by batch +# w/o this the trainer stores logits of all sample in memory... +# batch_eval_metrics: true + +per_device_train_batch_size: 16 +per_device_eval_batch_size: 16 +max_val_samples_per_ds: 1000 +# optim: schedule_free_adamw + +learning_rate: 0.00002 +# lr_scheduler_type: "constant_with_warmup" +neftune_noise_alpha: 1 +weight_decay: 0.01 +warmup_ratio: 0.1 + +# LoRA +lora_r: 8 +lora_dropout: 0.05 +target_modules: + - down_proj + - up_proj + - gate_proj + +# data +train_ds_names: +- fw_qa_tiny + +val_ds_names: +- fw_qa_tiny diff --git a/configs/pwc_hotpot_qa.yaml b/configs/pwc_hotpot_qa.yaml index 13eddd8..3da3f64 100644 --- a/configs/pwc_hotpot_qa.yaml +++ b/configs/pwc_hotpot_qa.yaml @@ -20,14 +20,15 @@ per_device_train_batch_size: 16 per_device_eval_batch_size: 16 max_val_samples_per_ds: 1000 # optim: schedule_free_adamw -learning_rate: 0.00001 + +learning_rate: 0.00002 # lr_scheduler_type: "constant_with_warmup" neftune_noise_alpha: 1 -weight_decay: 0.1 +weight_decay: 0.01 warmup_ratio: 0.1 # LoRA -lora_r: 16 +lora_r: 8 lora_dropout: 0.05 target_modules: - down_proj diff --git a/generate_fw_qa.py b/generate_fw_qa.py new file mode 100644 index 0000000..e709d13 --- /dev/null +++ b/generate_fw_qa.py @@ -0,0 +1,162 @@ +from glob import glob +import os +import re +import time +import random +import json +from numpy import mat +import requests +from dataclasses import dataclass, asdict +from typing import List, Optional + +import yaml +from openai import OpenAI +from datasets import load_dataset +from argparse import ArgumentParser, Namespace +from datetime import datetime +from tqdm import tqdm + + +from hyper_llm_modulator.utils import get_preprocessing_fn + + +# NOTE: Please, store your openai key with "export OPENAI_API_KEY=..." +api_key = os.environ.get("OPENAI_API_KEY") + + +SYSTEM_TEMPLATE = "You are a creative and helpful assistant." + +# based on Make Your LLM Fully Utilize the Context (https://arxiv.org/pdf/2404.16811) +PROMPT_TEMPLATE = ( + "Generate questions and corresponding answers from the given context. The questions should be highly specific to the " + "information provided in the context, not general questions that suits any context.\n\n" + "Rules to follow when generate the questions:\n" + "1. The questions must be fully answerable from information present in given context.\n" + "2. Make sure the questions are clear and unambiguous.\n" + "3. Phrases like 'based on the provided context', 'according to the context', etc, are not allowed to appear in " + "the questions.\n\n" + "Rules to follow when generate the answers:\n" + "1. The answers must use the information provided in the context.\n" + "2. Do not just copy words from the context. Answer the question in your own words.\n\n" + "Response with {n_qa_pairs} question-answer pairs. Use simple words and please be clear.\n" + "The question-answer pairs should be in the following format:\n" + "Question 1: {{question_1}}\n" + "Answer 1: {{answer_1}}\n" + "Question 2: {{question_2}}\n" + "Answer 2: {{answer_2}}\n" + "..." + "\n\n" + "### Context ###\n" + "{context}" +) + + +def get_prompt(context, n_qa_pairs): + prompt = PROMPT_TEMPLATE.format(context=context, n_qa_pairs=n_qa_pairs) + return prompt + + +def get_json_request(id, text, n_qa_pairs, gpt_model_name): + prompt = get_prompt(text, n_qa_pairs) + messages = [ + {"role": "system", "content": SYSTEM_TEMPLATE}, + {"role": "user", "content": prompt}, + ] + return { + "custom_id": f"{id}", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": gpt_model_name, + "messages": messages, + "temperature": 1.0, + "frequency_penalty": 0.2, + }, + } + + +def parse_args() -> Namespace: + p = ArgumentParser() + p.add_argument("--data_dir", type=str, default="tasks") + + # Every second, different unique id + current_time = datetime.now().strftime("%Y%m%d%H%M%S") + default_output_dir = f"generated_tasks/{current_time}" + p.add_argument("--generated_data_dir", type=str, default=default_output_dir) + p.add_argument("--gpt_model_name", type=str, default="gpt-4o-mini") + p.add_argument("--n_qa_pairs", type=int, default=2) + p.add_argument("--shard_pattern", type=str, default="*") + + return p.parse_args() + + +if __name__ == "__main__": + # Load the datasets + args = parse_args() + ds = load_dataset( + "parquet", + data_files=f"./data/raw_datasets/fineweb_sharded/{args.shard_pattern}.parquet", + split="train", + streaming=True, + ) + os.makedirs("openai_batches", exist_ok=True) + lines = [] + c = 0 + for i, sample in tqdm(enumerate(iter(ds))): + if len(lines) >= 20_000: + with open(f"openai_batches/fineweb_qa_pairs_{c}.jsonl", "w") as f: + for line in lines: + f.write(json.dumps(line) + "\n") + lines = [] + c += 1 + jsonl = get_json_request( + f"{c}_{i}", sample["text"], args.n_qa_pairs, args.gpt_model_name + ) + lines.append(jsonl) + + client = OpenAI() + + for file in glob("openai_batches/fineweb_qa_pairs_*.jsonl"): + print(f"Submitting batch {file}") + batch_input_file = client.files.create(file=open(file, "rb"), purpose="batch") + batch_input_file_id = batch_input_file.id + batch_info = client.batches.create( + input_file_id=batch_input_file_id, + endpoint="/v1/chat/completions", + completion_window="24h", + metadata={"description": f"{file}"}, + ) + print(f"Batch {file} submitted") + print(batch_info) + print("-" * 100) + + while ( + batch_info.completed_at is None + and batch_info.failed_at is None + and batch_info.expired_at is None + and batch_info.cancelled_at is None + ): + time.sleep(10) + batch_info = client.batches.retrieve(batch_info.id) + print(f"curtime: {datetime.now()}") + print(batch_info) + print() + + if batch_info.status == "completed": + print(f"Batch {file} completed") + res = client.files.content(batch_info.output_file_id) + + with open(f"{file.replace('.jsonl', '_res.jsonl')}", "w") as f: + for line in res.iter_lines(): + json_line = json.loads(line) + if json_line["error"]: + print(f"Error: {json_line['error']}") + print("Skipping line...") + continue + f.write(json.dumps(json_line) + "\n") + else: + print(f"Batch {file} failed") + print(batch_info) + print("-" * 100) + + time.sleep(10) diff --git a/generate_qa_parallel.py b/generate_qa_parallel.py new file mode 100644 index 0000000..9455195 --- /dev/null +++ b/generate_qa_parallel.py @@ -0,0 +1,493 @@ +# from https://github.com/openai/openai-cookbook/blob/main/examples/api_request_parallel_processor.py +""" +API REQUEST PARALLEL PROCESSOR + +Using the OpenAI API to process lots of text quickly takes some care. +If you trickle in a million API requests one by one, they'll take days to complete. +If you flood a million API requests in parallel, they'll exceed the rate limits and fail with errors. +To maximize throughput, parallel requests need to be throttled to stay under rate limits. + +This script parallelizes requests to the OpenAI API while throttling to stay under rate limits. + +Features: +- Streams requests from file, to avoid running out of memory for giant jobs +- Makes requests concurrently, to maximize throughput +- Throttles request and token usage, to stay under rate limits +- Retries failed requests up to {max_attempts} times, to avoid missing data +- Logs errors, to diagnose problems with requests + +Example command to call script: +``` +python examples/api_request_parallel_processor.py \ + --requests_filepath examples/data/example_requests_to_parallel_process.jsonl \ + --save_filepath examples/data/example_requests_to_parallel_process_results.jsonl \ + --request_url https://api.openai.com/v1/embeddings \ + --max_requests_per_minute 1500 \ + --max_tokens_per_minute 6250000 \ + --token_encoding_name cl100k_base \ + --max_attempts 5 \ + --logging_level 20 +``` + +Inputs: +- requests_filepath : str + - path to the file containing the requests to be processed + - file should be a jsonl file, where each line is a json object with API parameters and an optional metadata field + - e.g., {"model": "text-embedding-3-small", "input": "embed me", "metadata": {"row_id": 1}} + - as with all jsonl files, take care that newlines in the content are properly escaped (json.dumps does this automatically) + - an example file is provided at examples/data/example_requests_to_parallel_process.jsonl + - the code to generate the example file is appended to the bottom of this script +- save_filepath : str, optional + - path to the file where the results will be saved + - file will be a jsonl file, where each line is an array with the original request plus the API response + - e.g., [{"model": "text-embedding-3-small", "input": "embed me"}, {...}] + - if omitted, results will be saved to {requests_filename}_results.jsonl +- request_url : str, optional + - URL of the API endpoint to call + - if omitted, will default to "https://api.openai.com/v1/chat/completions" +- api_key : str, optional + - API key to use + - if omitted, the script will attempt to read it from an environment variable {os.getenv("OPENAI_API_KEY")} +- max_requests_per_minute : float, optional + - target number of requests to make per minute (will make less if limited by tokens) + - leave headroom by setting this to 50% or 75% of your limit + - if requests are limiting you, try batching multiple embeddings or completions into one request + - if omitted, will default to 1,500 +- max_tokens_per_minute : float, optional + - target number of tokens to use per minute (will use less if limited by requests) + - leave headroom by setting this to 50% or 75% of your limit + - if omitted, will default to 125,000 +- token_encoding_name : str, optional + - name of the token encoding used, as defined in the `tiktoken` package + - if omitted, will default to "o200k_base" (used by `gpt-4o-mini`) + - see https://cookbook.openai.com/examples/how_to_count_tokens_with_tiktoken +- max_attempts : int, optional + - number of times to retry a failed request before giving up + - if omitted, will default to 5 +- logging_level : int, optional + - level of logging to use; higher numbers will log fewer messages + - 40 = ERROR; will log only when requests fail after all retries + - 30 = WARNING; will log when requests his rate limits or other errors + - 20 = INFO; will log when requests start and the status at finish + - 10 = DEBUG; will log various things as the loop runs to see when they occur + - if omitted, will default to 20 (INFO). + +The script is structured as follows: + - Imports + - Define main() + - Initialize things + - In main loop: + - Get next request if one is not already waiting for capacity + - Update available token & request capacity + - If enough capacity available, call API + - The loop pauses if a rate limit error is hit + - The loop breaks when no tasks remain + - Define dataclasses + - StatusTracker (stores script metadata counters; only one instance is created) + - APIRequest (stores API inputs, outputs, metadata; one method to call API) + - Define functions + - api_endpoint_from_url (extracts API endpoint from request URL) + - append_to_jsonl (writes to results file) + - num_tokens_consumed_from_request (bigger function to infer token usage from request) + - task_id_generator_function (yields 0, 1, 2, ...) + - Run main() +""" + +# imports +import aiohttp # for making API calls concurrently +import argparse # for running script from command line +import asyncio # for running API calls concurrently +import json # for saving results to a jsonl file +import logging # for logging rate limit warnings and other messages +import os # for reading API key +import re # for matching endpoint from request URL +import tiktoken # for counting tokens +import time # for sleeping after rate limit is hit +from dataclasses import ( + dataclass, + field, +) # for storing API inputs, outputs, and metadata + + +async def process_api_requests_from_file( + requests_filepath: str, + save_filepath: str, + request_url: str, + api_key: str, + max_requests_per_minute: float, + max_tokens_per_minute: float, + token_encoding_name: str, + max_attempts: int, + logging_level: int, +): + """Processes API requests in parallel, throttling to stay under rate limits.""" + # constants + seconds_to_pause_after_rate_limit_error = 15 + seconds_to_sleep_each_loop = ( + 0.01 # 10 ms limits max throughput to 100 requests per second + ) + + # initialize logging + logging.basicConfig(level=logging_level) + logging.debug(f"Logging initialized at level {logging_level}") + + # infer API endpoint and construct request header + api_endpoint = api_endpoint_from_url(request_url) + request_header = {"Authorization": f"Bearer {api_key}"} + # use api-key header for Azure deployments + if "/deployments" in request_url: + request_header = {"api-key": f"{api_key}"} + + # initialize trackers + queue_of_requests_to_retry = asyncio.Queue() + task_id_generator = ( + task_id_generator_function() + ) # generates integer IDs of 0, 1, 2, ... + status_tracker = ( + StatusTracker() + ) # single instance to track a collection of variables + next_request = None # variable to hold the next request to call + + # initialize available capacity counts + available_request_capacity = max_requests_per_minute + available_token_capacity = max_tokens_per_minute + last_update_time = time.time() + + # initialize flags + file_not_finished = True # after file is empty, we'll skip reading it + logging.debug(f"Initialization complete.") + + # initialize file reading + with open(requests_filepath) as file: + # `requests` will provide requests one at a time + requests = file.__iter__() + logging.debug(f"File opened. Entering main loop") + async with aiohttp.ClientSession() as session: # Initialize ClientSession here + while True: + # get next request (if one is not already waiting for capacity) + if next_request is None: + if not queue_of_requests_to_retry.empty(): + next_request = queue_of_requests_to_retry.get_nowait() + logging.debug( + f"Retrying request {next_request.task_id}: {next_request}" + ) + elif file_not_finished: + try: + # get new request + request_json = json.loads(next(requests)) + next_request = APIRequest( + task_id=next(task_id_generator), + request_json=request_json, + token_consumption=num_tokens_consumed_from_request( + request_json, api_endpoint, token_encoding_name + ), + attempts_left=max_attempts, + metadata=request_json.pop("metadata", None), + ) + status_tracker.num_tasks_started += 1 + status_tracker.num_tasks_in_progress += 1 + logging.debug( + f"Reading request {next_request.task_id}: {next_request}" + ) + except StopIteration: + # if file runs out, set flag to stop reading it + logging.debug("Read file exhausted") + file_not_finished = False + + # update available capacity + current_time = time.time() + seconds_since_update = current_time - last_update_time + available_request_capacity = min( + available_request_capacity + + max_requests_per_minute * seconds_since_update / 60.0, + max_requests_per_minute, + ) + available_token_capacity = min( + available_token_capacity + + max_tokens_per_minute * seconds_since_update / 60.0, + max_tokens_per_minute, + ) + last_update_time = current_time + + # if enough capacity available, call API + if next_request: + next_request_tokens = next_request.token_consumption + if ( + available_request_capacity >= 1 + and available_token_capacity >= next_request_tokens + ): + # update counters + available_request_capacity -= 1 + available_token_capacity -= next_request_tokens + next_request.attempts_left -= 1 + + # call API + asyncio.create_task( + next_request.call_api( + session=session, + request_url=request_url, + request_header=request_header, + retry_queue=queue_of_requests_to_retry, + save_filepath=save_filepath, + status_tracker=status_tracker, + ) + ) + next_request = None # reset next_request to empty + + # if all tasks are finished, break + if status_tracker.num_tasks_in_progress == 0: + break + + # main loop sleeps briefly so concurrent tasks can run + await asyncio.sleep(seconds_to_sleep_each_loop) + + # if a rate limit error was hit recently, pause to cool down + seconds_since_rate_limit_error = ( + time.time() - status_tracker.time_of_last_rate_limit_error + ) + if ( + seconds_since_rate_limit_error + < seconds_to_pause_after_rate_limit_error + ): + remaining_seconds_to_pause = ( + seconds_to_pause_after_rate_limit_error + - seconds_since_rate_limit_error + ) + await asyncio.sleep(remaining_seconds_to_pause) + # ^e.g., if pause is 15 seconds and final limit was hit 5 seconds ago + logging.warn( + f"Pausing to cool down until {time.ctime(status_tracker.time_of_last_rate_limit_error + seconds_to_pause_after_rate_limit_error)}" + ) + + # after finishing, log final status + logging.info( + f"""Parallel processing complete. Results saved to {save_filepath}""" + ) + if status_tracker.num_tasks_failed > 0: + logging.warning( + f"{status_tracker.num_tasks_failed} / {status_tracker.num_tasks_started} requests failed. Errors logged to {save_filepath}." + ) + if status_tracker.num_rate_limit_errors > 0: + logging.warning( + f"{status_tracker.num_rate_limit_errors} rate limit errors received. Consider running at a lower rate." + ) + + +# dataclasses + + +@dataclass +class StatusTracker: + """Stores metadata about the script's progress. Only one instance is created.""" + + num_tasks_started: int = 0 + num_tasks_in_progress: int = 0 # script ends when this reaches 0 + num_tasks_succeeded: int = 0 + num_tasks_failed: int = 0 + num_rate_limit_errors: int = 0 + num_api_errors: int = 0 # excluding rate limit errors, counted above + num_other_errors: int = 0 + time_of_last_rate_limit_error: int = 0 # used to cool off after hitting rate limits + + +@dataclass +class APIRequest: + """Stores an API request's inputs, outputs, and other metadata. Contains a method to make an API call.""" + + task_id: int + request_json: dict + token_consumption: int + attempts_left: int + metadata: dict + result: list = field(default_factory=list) + + async def call_api( + self, + session: aiohttp.ClientSession, + request_url: str, + request_header: dict, + retry_queue: asyncio.Queue, + save_filepath: str, + status_tracker: StatusTracker, + ): + """Calls the OpenAI API and saves results.""" + logging.info(f"Starting request #{self.task_id}") + error = None + try: + async with session.post( + url=request_url, headers=request_header, json=self.request_json + ) as response: + response = await response.json() + if "error" in response: + logging.warning( + f"Request {self.task_id} failed with error {response['error']}" + ) + status_tracker.num_api_errors += 1 + error = response + if "rate limit" in response["error"].get("message", "").lower(): + status_tracker.time_of_last_rate_limit_error = time.time() + status_tracker.num_rate_limit_errors += 1 + status_tracker.num_api_errors -= ( + 1 # rate limit errors are counted separately + ) + + except ( + Exception + ) as e: # catching naked exceptions is bad practice, but in this case we'll log & save them + logging.warning(f"Request {self.task_id} failed with Exception {e}") + status_tracker.num_other_errors += 1 + error = e + if error: + self.result.append(error) + if self.attempts_left: + retry_queue.put_nowait(self) + else: + logging.error( + f"Request {self.request_json} failed after all attempts. Saving errors: {self.result}" + ) + data = ( + [self.request_json, [str(e) for e in self.result], self.metadata] + if self.metadata + else [self.request_json, [str(e) for e in self.result]] + ) + append_to_jsonl(data, save_filepath) + status_tracker.num_tasks_in_progress -= 1 + status_tracker.num_tasks_failed += 1 + else: + data = ( + [self.request_json, response, self.metadata] + if self.metadata + else [self.request_json, response] + ) + append_to_jsonl(data, save_filepath) + status_tracker.num_tasks_in_progress -= 1 + status_tracker.num_tasks_succeeded += 1 + logging.debug(f"Request {self.task_id} saved to {save_filepath}") + + +# functions + + +def api_endpoint_from_url(request_url): + """Extract the API endpoint from the request URL.""" + match = re.search("^https://[^/]+/v\\d+/(.+)$", request_url) + if match is None: + # for Azure OpenAI deployment urls + match = re.search( + r"^https://[^/]+/openai/deployments/[^/]+/(.+?)(\?|$)", request_url + ) + return match[1] + + +def append_to_jsonl(data, filename: str) -> None: + """Append a json payload to the end of a jsonl file.""" + json_string = json.dumps(data) + with open(filename, "a") as f: + f.write(json_string + "\n") + + +def num_tokens_consumed_from_request( + request_json: dict, + api_endpoint: str, + token_encoding_name: str, +): + """Count the number of tokens in the request. Only supports completion and embedding requests.""" + encoding = tiktoken.get_encoding(token_encoding_name) + # if completions request, tokens = prompt + n * max_tokens + if api_endpoint.endswith("completions"): + max_tokens = request_json.get("max_tokens", 15) + n = request_json.get("n", 1) + completion_tokens = n * max_tokens + + # chat completions + if api_endpoint.startswith("chat/"): + num_tokens = 0 + for message in request_json["messages"]: + num_tokens += ( + 4 # every message follows {role/name}\n{content}\n + ) + for key, value in message.items(): + num_tokens += len(encoding.encode(value)) + if key == "name": # if there's a name, the role is omitted + num_tokens -= 1 # role is always required and always 1 token + num_tokens += 2 # every reply is primed with assistant + return num_tokens + completion_tokens + # normal completions + else: + prompt = request_json["prompt"] + if isinstance(prompt, str): # single prompt + prompt_tokens = len(encoding.encode(prompt)) + num_tokens = prompt_tokens + completion_tokens + return num_tokens + elif isinstance(prompt, list): # multiple prompts + prompt_tokens = sum([len(encoding.encode(p)) for p in prompt]) + num_tokens = prompt_tokens + completion_tokens * len(prompt) + return num_tokens + else: + raise TypeError( + 'Expecting either string or list of strings for "prompt" field in completion request' + ) + # if embeddings request, tokens = input tokens + elif api_endpoint == "embeddings": + input = request_json["input"] + if isinstance(input, str): # single input + num_tokens = len(encoding.encode(input)) + return num_tokens + elif isinstance(input, list): # multiple inputs + num_tokens = sum([len(encoding.encode(i)) for i in input]) + return num_tokens + else: + raise TypeError( + 'Expecting either string or list of strings for "inputs" field in embedding request' + ) + # more logic needed to support other API calls (e.g., edits, inserts, DALL-E) + else: + raise NotImplementedError( + f'API endpoint "{api_endpoint}" not implemented in this script' + ) + + +def task_id_generator_function(): + """Generate integers 0, 1, 2, and so on.""" + task_id = 0 + while True: + yield task_id + task_id += 1 + + +# run script + + +if __name__ == "__main__": + # parse command line arguments + parser = argparse.ArgumentParser() + parser.add_argument("--requests_filepath") + parser.add_argument("--save_filepath", default=None) + parser.add_argument( + "--request_url", default="https://api.openai.com/v1/chat/completions" + ) + parser.add_argument("--api_key", default=os.getenv("OPENAI_API_KEY")) + parser.add_argument("--max_requests_per_minute", type=int, default=5_000 * 0.5) + parser.add_argument("--max_tokens_per_minute", type=int, default=2_000_000 * 0.8) + parser.add_argument("--token_encoding_name", default="o200k_base") + parser.add_argument("--max_attempts", type=int, default=5) + parser.add_argument("--logging_level", default=logging.INFO) + args = parser.parse_args() + + if args.save_filepath is None: + args.save_filepath = args.requests_filepath.replace(".jsonl", "_res.jsonl") + + # run script + asyncio.run( + process_api_requests_from_file( + requests_filepath=args.requests_filepath, + save_filepath=args.save_filepath, + request_url=args.request_url, + api_key=args.api_key, + max_requests_per_minute=float(args.max_requests_per_minute), + max_tokens_per_minute=float(args.max_tokens_per_minute), + token_encoding_name=args.token_encoding_name, + max_attempts=int(args.max_attempts), + logging_level=int(args.logging_level), + ) + ) diff --git a/hyperlora/configs.py b/hyperlora/configs.py index 05d18c0..3cbc430 100644 --- a/hyperlora/configs.py +++ b/hyperlora/configs.py @@ -81,7 +81,9 @@ class ArgumentParser(HfArgumentParser): obj = data_class(**inputs) outputs.append(obj) - + for arg in other_args: + if arg not in used_args: + raise ValueError(f"Argument provided not found in dataclass: {arg}") return outputs def parse(self) -> DataClassType | tuple[DataClassType]: @@ -122,10 +124,36 @@ class ExperimentSetup(str, Enum): @dataclass class TrainingArguments(TrainingArguments): + dataloader_pin_memory: bool = field( + default=True, + metadata={"help": "Whether to pin memory in data loaders or not."}, + ) + dataloader_persistent_workers: bool = field( + default=True, + metadata={ + "help": "Whether to keep the workers alive after a dataset has been consumed once." + }, + ) + dataloader_prefetch_factor: int = field( + default=2, + metadata={"help": "Number of batches loaded in advance by each worker."}, + ) + dataloader_num_workers: int = field( + default=4, + metadata={"help": "Number of subprocesses to use for data loading."}, + ) optim: str = field( default="adamw_torch_fused", metadata={"help": "Optimizer."}, ) + adam_beta1: float = field( + default=0.9, + metadata={"help": "Adam beta 1."}, + ) + adam_beta2: float = field( + default=0.95, + metadata={"help": "Adam beta 2."}, + ) eval_on_start: bool = field( default=True, metadata={"help": "Whether to evaluate on the start of training."}, diff --git a/hyperlora/data_utils.py b/hyperlora/data_utils.py index d61640c..0061725 100644 --- a/hyperlora/data_utils.py +++ b/hyperlora/data_utils.py @@ -2,7 +2,7 @@ import logging import numpy as np from typing import Any, Callable, Iterator, Optional -from datasets import load_dataset +from datasets import load_dataset, IterableDataset from training_utils import TRAINING_TASK from transformers import PreTrainedTokenizerBase @@ -35,22 +35,16 @@ DS_KWARGS = { validation=dict(path="sggetao/PwC", split="train[:1000]"), test=dict(path="sggetao/PwC", split="test"), ), - "fineweb_tiny": dict( + "fw_qa_tiny": dict( train=dict( path="parquet", - data_files="data/raw_datasets/fineweb_sharded/00000.parquet", - split="train[2000:]", - streaming=True, + data_files="data/raw_datasets/fw_qa/00000.parquet", + split="train", ), validation=dict( path="parquet", - data_files="data/raw_datasets/fineweb_sharded/00000.parquet", - split="train[1000:2000]", - ), - test=dict( - path="parquet", - data_files="data/raw_datasets/fineweb_sharded/00000.parquet", - split="train[:1000]", + data_files="data/raw_datasets/fw_qa/00000_val.parquet", + split="train", ), ), } @@ -58,11 +52,12 @@ DS_KWARGS = { def get_ds_kwargs(ds_name: str, split: str) -> dict[str, Any]: if (ds_name not in DS_KWARGS) or (split not in DS_KWARGS[ds_name]): + kwargs = dict(path=ds_name, split=split) logger.warning( f"No dataset kwargs found for '{ds_name}' with split '{split}'.\n" - f"Using default kwargs: path={ds_name}, split={split}" + f"Using default kwargs: {kwargs}" ) - return dict(split=split) + return kwargs return DS_KWARGS[ds_name][split] @@ -191,15 +186,27 @@ def get_tokenized_dataset( f"Failed to load dataset {ds_name} with split {split}. Error: {e}\nSkipping..." ) return None - ds = ds.map(get_preprocessing_fn(ds_name), remove_columns=ds.column_names) + cols_to_remove = [ + col for col in ds.column_names if col not in ["context", "prompt", "response"] + ] + ds = ds.map(get_preprocessing_fn(ds_name)) + ds = ds.remove_columns(cols_to_remove) ds = ds.filter(filter_none, batched=True) ds = ds.filter(filter_long_samples, batched=True) - if split != "test": + if split == "train": if add_negative_prompt: ds = ds.map(add_negative_prompt_fn, batched=True, batch_size=None) if add_repeat_prompt and "context_numbers" not in ds_name: ds = ds.map(add_repeat_prompt_fn, batched=True, batch_size=None) + tokenized_ds = construct_and_tokenize_ctx_qa( + tokenizer, tokenizer_kwargs, add_ctx_to_chat, use_kl_loss, need_ctx_ids, ds + ) + return tokenized_ds + +def construct_and_tokenize_ctx_qa( + tokenizer, tokenizer_kwargs, add_ctx_to_chat, use_kl_loss, need_ctx_ids, ds +): # for sft + chat_model, we need to convert the dataset to chat format # add "messages" field ds = ds.map( @@ -220,12 +227,6 @@ def get_tokenized_dataset( }, ) - other_cols = [ - col - for col in tokenized_ds.column_names - if col not in ["input_ids", "attention_mask", "labels"] - ] - # for use_kl_loss, we need "chat_ids" and "chat_attn_mask" if use_kl_loss: tokenized_ds = tokenized_ds.map( @@ -245,6 +246,7 @@ def get_tokenized_dataset( "tokenizer_kwargs": tokenizer_kwargs, "for_kl_loss": True, }, + num_proc=16, ) if need_ctx_ids: @@ -254,9 +256,13 @@ def get_tokenized_dataset( tokenized_ds = tokenized_ds.map( tokenize_ctx_text, fn_kwargs={"tokenizer": tokenizer}, + batched=True, + num_proc=16, ) - tokenized_ds = tokenized_ds.remove_columns(other_cols) + tokenized_ds = tokenized_ds.remove_columns( + ["messages", "chat", "context", "prompt", "response"] + ) tokenized_ds.set_format(type="pt") validate_columns(tokenized_ds) return tokenized_ds diff --git a/hyperlora/intx_sft.py b/hyperlora/intx_sft.py old mode 100644 new mode 100755 index 72e9e5f..cddd269 --- a/hyperlora/intx_sft.py +++ b/hyperlora/intx_sft.py @@ -22,7 +22,13 @@ from data_utils import ( tokenize_chat_messages, tokenize_ctx_text, ) -from datasets import concatenate_datasets, disable_caching, load_dataset +from datasets import ( + concatenate_datasets, + interleave_datasets, + disable_caching, + load_dataset, + IterableDataset, +) from model_loading import get_lora_config, get_model_and_tokenizer from modeling_utils import ( EarlyExit, @@ -40,6 +46,7 @@ from transformers import ( HfArgumentParser, set_seed, ) +from torch.utils.data import DataLoader from utils import ( extract_cli_args, get_base_model, @@ -292,6 +299,7 @@ def main(): add_repeat_prompt=ctx_args.add_repeat_prompt, add_negative_prompt=ctx_args.add_negative_prompt, use_kl_loss=ctx_args.use_kl_loss, + # streaming=data_args.streaming, ) tokenized_ds = {} for split, ds_names in zip( @@ -316,22 +324,28 @@ def main(): train_ds[ds_name] = train_ds[ds_name].skip(n_val_samples) val_ds[ds_name] = ds - val_ds_size = len(val_ds[ds_name]) - val_indices = np.random.permutation(val_ds_size)[:n_val_samples] + val_indices = np.random.permutation(len(ds))[:n_val_samples] val_ds[ds_name] = val_ds[ds_name].select(val_indices) - train_ds = concatenate_datasets(list(train_ds.values())) + # train_ds = concatenate_datasets(list(train_ds.values())) + + # total_len = sum(len(ds) for ds in train_ds.values()) + train_ds_len = [len(ds) for ds in train_ds.values()] + total_len = sum(train_ds_len) + train_ds = interleave_datasets( + list(train_ds.values()), + probabilities=[l / total_len for l in train_ds_len], + ) + val_train_indices = np.random.permutation(len(train_ds))[:500] val_ds["train"] = train_ds.select(val_train_indices) test_ds = dict() if "test" in tokenized_ds: + n_test_samples = data_args.max_test_samples_per_ds for ds_name, ds in tokenized_ds["test"].items(): test_ds[ds_name] = ds - test_ds_size = len(test_ds[ds_name]) - test_indices = np.random.permutation(test_ds_size)[ - : data_args.max_test_samples_per_ds - ] + test_indices = np.random.permutation(len(ds))[:n_test_samples] test_ds[ds_name] = test_ds[ds_name].select(test_indices) logger.info(f"train_ds: {train_ds}") @@ -471,7 +485,7 @@ def main(): test_ds, partial(train_collator, tokenizer=tokenizer), partial(generation_collator, tokenizer=tokenizer), - partial( + compute_metrics=partial( compute_metrics, evaluator=Evaluator( [compute_per_token_acc, compute_prefix_matching, compute_perplexity] @@ -484,6 +498,7 @@ def main(): if __name__ == "__main__": + os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true" os.environ["TOKENIZERS_PARALLELISM"] = "true" os.environ["WANDB_PROJECT"] = "ctx_to_lora" os.environ["WANDB_WATCH"] = "" # "all" diff --git a/install.sh b/install.sh deleted file mode 100755 index 48263df..0000000 --- a/install.sh +++ /dev/null @@ -1,14 +0,0 @@ -# install unsloth -# see https://docs.unsloth.ai/get-started/install-update/conda-install -conda create --name unsloth_env \ - python=3.10 \ - pytorch-cuda=<11.8/12.1> \ - pytorch cudatoolkit xformers -c pytorch -c nvidia -c xformers \ - -y -conda activate unsloth_env - -pip install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git" - -pip install --no-deps "trl<0.9.0" peft accelerate bitsandbytes - -pip install -r requirements.txt diff --git a/post_process_fw_qa.py b/post_process_fw_qa.py new file mode 100644 index 0000000..704b776 --- /dev/null +++ b/post_process_fw_qa.py @@ -0,0 +1,216 @@ +import re +import json +import random +from glob import glob + +import pandas as pd +import numpy as np +from datasets import Dataset + + +def get_repeat_prompts(ds): + unique_contexts = set() + ctxs, prompts, responses = [], [], [] + for ctx, prompt, response in zip(ds["context"], ds["prompt"], ds["response"]): + # Only process if the context is not already in the set + if ctx in unique_contexts: + continue + + unique_contexts.add(ctx) + ctxs.append(ctx) + responses.append(ctx) + prompts.append("Repeat the text above.") + + print(f"Adding repeat prompt...") + print(f"# unique contexts: {len(unique_contexts)}") + + return Dataset.from_pandas( + pd.DataFrame( + dict( + context=ctxs, + response=responses, + prompt=prompts, + ) + ) + ) + + +def get_negative_prompts(ds): + unique_contexts = set() + ctxs, prompts, responses = [], [], [] + keywords = [ + "repeat", + "rephrase", + "summarize", + "rewrite", + "title", + "keyword", + "continuation", + ] + + for ctx, prompt, response in zip(ds["context"], ds["prompt"], ds["response"]): + if ctx in unique_contexts: + continue + if any(keyword in prompt for keyword in keywords): + # Skip samples where the prompt contains any of the specified keywords + continue + + unique_contexts.add(ctx) + ctxs.append(ctx) + prompts.append(prompt) + responses.append(response) + + print(f"Adding negative prompt...") + print(f"# unique contexts: {len(unique_contexts)}") + + # remove one last sample if the number of samples is odd + if len(ctxs) % 2 != 0: + ctxs.pop() + prompts.pop() + responses.pop() + + # to make sure that the negative prompt/response is not the same as the original + indices = list(np.random.permutation(len(ctxs))) + list( + np.random.permutation(len(ctxs)) + ) + neg_ctxs, neg_prompts, neg_responses = [], [], [] + for idx in range(0, len(indices), 2): + i = indices[idx] + j = indices[idx + 1] + neg_ctxs.append(ctxs[i]) + neg_prompts.append(ctxs[j] + "\n\n" + prompts[j]) + neg_responses.append(responses[j]) + + return Dataset.from_pandas( + pd.DataFrame( + dict( + context=neg_ctxs, + prompt=neg_prompts, + response=neg_responses, + ) + ) + ) + + +def postprocess_qa_pairs(res_txt: str): + """ + Postprocesses the QA pairs from the response text. + + Args: + res_txt: The response text. + n_qa_pairs: The number of QA pairs. + + Returns: + A tuple of two lists, the first containing the questions and the second containing the answers. + """ + # capture everything after each "Question {number}:" until "Answer" + q_pattern = r"Question \d+:(.*?)(?=Answer|$)" # thanks chatgpt + questions = re.findall(q_pattern, res_txt, flags=re.S) + + a_pattern = r"Answer \d+:(.*?)(?=Question|$)" # thanks chatgpt + answers = re.findall(a_pattern, res_txt, flags=re.S) + + if len(questions) != len(answers): + print(f"Warning---number of questions and answers do not match") + print(f"Number of questions: {len(questions)}") + print(f"Number of answers: {len(answers)}") + + out_q = [] + out_a = [] + if (len(questions) > 0) and (len(answers) > 0): + for i in range(min(len(questions), len(answers))): + out_q.append(questions[i].strip()) + out_a.append(answers[i].strip()) + + return out_q, out_a + + +def get_response_txt(res_item): + try: + return res_item["response"]["body"]["choices"][0]["message"]["content"] + except Exception as e: + print(f"Error getting response: {e}") + return None + + +def get_prompt_txt(prompt_item): + try: + return prompt_item["body"]["messages"][1]["content"] + except Exception as e: + print(f"Error getting context: {e}") + return None + + +if __name__ == "__main__": + res_files = glob("openai_batches/fineweb_qa_pairs_*_res.jsonl") + prompt_files = [f.replace("_res.jsonl", ".jsonl") for f in res_files] + + for res_file, prompt_file in zip(res_files, prompt_files): + print(f"Processing {res_file} and {prompt_file}") + idx = int(prompt_file.split("/")[-1].split("_")[-1].split(".jsonl")[0]) + res_data = [] + prompt_data = [] + samples = [] + + # Load data from res_file + with open(res_file, "r") as f: + for line in f: + res_data.append(json.loads(line)) + + # Load data from prompt_file and create an index + prompt_data_index = {} + with open(prompt_file, "r") as f: + for line in f: + prompt_item = json.loads(line) + prompt_data_index[prompt_item.get("custom_id")] = prompt_item + + # Match JSON objects based on custom_id using the index + for res_item in res_data: + custom_id = res_item.get("custom_id") + prompt_item = prompt_data_index.get(custom_id) + if not prompt_item: + print(f"No prompt item found for {custom_id}") + continue + res_txt = get_response_txt(res_item) + if not res_txt: + print(f"No response text found for {custom_id}") + continue + questions, answers = postprocess_qa_pairs(res_txt) + if not questions or not answers: + print(f"No questions or answers found for {custom_id}") + continue + context = get_prompt_txt(prompt_item).split("### Context ###")[-1] + for q, a in zip(questions, answers): + samples.append( + { + "context": context, + "prompt": q, + "response": a, + } + ) + + print(f"Found {len(samples)} samples in {res_file} and {prompt_file}") + random.shuffle(samples) + df = pd.DataFrame(samples) + ds = Dataset.from_pandas(df) + val_ds = ds.take(100) + ds = ds.skip(100) + + ds.to_parquet(f"data/raw_datasets/fw_qa/{idx:05d}.parquet") + val_ds.to_parquet(f"data/raw_datasets/fw_qa/{idx:05d}_val.parquet") + print(f"Saved to data/raw_datasets/fw_qa/{idx:05d}.parquet") + print(f"Saved to data/raw_datasets/fw_qa/{idx:05d}_val.parquet") + + # add repeat and negative prompt + repeat_ds = get_repeat_prompts(ds).shuffle(seed=42) + repeat_ds.to_parquet(f"data/raw_datasets/fw_qa/{idx:05d}_repeat.parquet") + print(f"repeat_ds: {len(repeat_ds)}") + print( + f"Saved {idx:05d}_repeat.parquet to data/raw_datasets/fw_qa/{idx:05d}_repeat.parquet" + ) + neg_ds = get_negative_prompts(ds).shuffle(seed=42) + neg_ds.to_parquet(f"data/raw_datasets/fw_qa/{idx:05d}_neg.parquet") + print(f"neg_ds: {len(neg_ds)}") + print( + f"Saved {idx:05d}_neg.parquet to data/raw_datasets/fw_qa/{idx:05d}_neg.parquet" + ) diff --git a/post_process_parallel_qa.py b/post_process_parallel_qa.py new file mode 100644 index 0000000..9ea7641 --- /dev/null +++ b/post_process_parallel_qa.py @@ -0,0 +1,206 @@ +import json +import re +import os +from glob import glob + +import numpy as np +import pandas as pd +from datasets import Dataset + + +def get_repeat_prompts(ds): + unique_contexts = set() + ctxs, prompts, responses = [], [], [] + for ctx, prompt, response in zip(ds["context"], ds["prompt"], ds["response"]): + # Only process if the context is not already in the set + if ctx in unique_contexts: + continue + + unique_contexts.add(ctx) + ctxs.append(ctx) + responses.append(ctx) + prompts.append("Repeat the text above.") + + print(f"Adding repeat prompt...") + print(f"# unique contexts: {len(unique_contexts)}") + + return Dataset.from_pandas( + pd.DataFrame( + dict( + context=ctxs, + response=responses, + prompt=prompts, + ) + ) + ) + + +def get_negative_prompts(ds): + unique_contexts = set() + ctxs, prompts, responses = [], [], [] + keywords = [ + "repeat", + "rephrase", + "summarize", + "rewrite", + "title", + "keyword", + "continuation", + ] + + for ctx, prompt, response in zip(ds["context"], ds["prompt"], ds["response"]): + if ctx in unique_contexts: + continue + if any(keyword in prompt for keyword in keywords): + # Skip samples where the prompt contains any of the specified keywords + continue + + unique_contexts.add(ctx) + ctxs.append(ctx) + prompts.append(prompt) + responses.append(response) + + print(f"Adding negative prompt...") + print(f"# unique contexts: {len(unique_contexts)}") + + # remove one last sample if the number of samples is odd + if len(ctxs) % 2 != 0: + ctxs.pop() + prompts.pop() + responses.pop() + + # to make sure that the negative prompt/response is not the same as the original + indices = list(np.random.permutation(len(ctxs))) + list( + np.random.permutation(len(ctxs)) + ) + neg_ctxs, neg_prompts, neg_responses = [], [], [] + for idx in range(0, len(indices), 2): + i = indices[idx] + j = indices[idx + 1] + neg_ctxs.append(ctxs[i]) + neg_prompts.append(ctxs[j] + "\n\n" + prompts[j]) + neg_responses.append(responses[j]) + + return Dataset.from_pandas( + pd.DataFrame( + dict( + context=neg_ctxs, + prompt=neg_prompts, + response=neg_responses, + ) + ) + ) + + +def postprocess_qa_pairs(res_txt: str): + """ + Postprocesses the QA pairs from the response text. + + Args: + res_txt: The response text. + n_qa_pairs: The number of QA pairs. + + Returns: + A tuple of two lists, the first containing the questions and the second containing the answers. + """ + # capture everything after each "Question {number}:" until "Answer" + q_pattern = r"Question \d+:(.*?)(?=Answer|$)" # thanks chatgpt + questions = re.findall(q_pattern, res_txt, flags=re.S) + + a_pattern = r"Answer \d+:(.*?)(?=Question|$)" # thanks chatgpt + answers = re.findall(a_pattern, res_txt, flags=re.S) + + c_pattern = r"Continuation:\s*(.*)" # thanks chatgpt + continuations = re.findall(c_pattern, res_txt, flags=re.S) + + if len(questions) != len(answers): + print(f"Warning---number of questions and answers do not match") + print(f"Number of questions: {len(questions)}") + print(f"Number of answers: {len(answers)}") + + out_q = [] + out_a = [] + if (len(questions) > 0) and (len(answers) > 0): + for i in range(min(len(questions), len(answers))): + out_q.append(questions[i].strip()) + out_a.append(answers[i].strip()) + + if len(continuations) > 0: + assert len(continuations) == 1 + continuations[0] = continuations[0].strip() + + return out_q, out_a, continuations + + +if __name__ == "__main__": + file_pattern = os.sys.argv[1] + files = glob(file_pattern) + print(f"Files: {files}") + print(f"Processing {len(files)} files") + for file in files: + fname = os.path.basename(file).replace(".jsonl", "") + samples = [] + unique_contexts = set() + with open(file, "r") as f: + while True: + line = f.readline() + if not line: + break # End of file + try: + json_obj = json.loads(line) + except Exception as e: + print(f"Error loading JSON: {e}") + continue + + if not isinstance(json_obj[1], dict): + print(f"Error: {json_obj[1]} is not a dictionary") + continue + + prompt_txt = json_obj[0]["messages"][1]["content"].split( + "### Context ###" + )[-1] + res_txt = json_obj[1]["choices"][0]["message"]["content"] + questions, answers, continuations = postprocess_qa_pairs(res_txt) + unique_contexts.add(prompt_txt) + for q, a in zip(questions, answers): + samples.append( + { + "context": prompt_txt, + "prompt": q, + "response": a, + } + ) + if continuations: + samples.append( + { + "context": prompt_txt, + "prompt": "Write a 1-paragraph continuation of the text.", + "response": continuations[0], + } + ) + df = pd.DataFrame(samples) + + # raise NotImplementedError + ds = Dataset.from_pandas(df).shuffle(seed=42) + print( + f"Created {len(ds)} samples from {file} with {len(unique_contexts)} unique contexts" + ) + val_ds = ds.take(500) + ds = ds.skip(500) + ds.to_parquet(f"data/raw_datasets/ctx_qa/{fname}.parquet") + val_ds.to_parquet(f"data/raw_datasets/ctx_qa/{fname}_val.parquet") + print(f"Saved {fname} to data/raw_datasets/ctx_qa/{fname}.parquet") + print(f"Saved {fname}_val to data/raw_datasets/ctx_qa/{fname}_val.parquet") + # add repeat and negative prompt + repeat_ds = get_repeat_prompts(ds).shuffle(seed=42) + repeat_ds.to_parquet(f"data/raw_datasets/ctx_qa/{fname}_repeat.parquet") + print(f"repeat_ds: {len(repeat_ds)}") + print( + f"Saved {fname}_repeat.parquet to data/raw_datasets/ctx_qa/{fname}_repeat.parquet" + ) + neg_ds = get_negative_prompts(ds).shuffle(seed=42) + neg_ds.to_parquet(f"data/raw_datasets/ctx_qa/{fname}_neg.parquet") + print(f"neg_ds: {len(neg_ds)}") + print( + f"Saved {fname}_neg.parquet to data/raw_datasets/ctx_qa/{fname}_neg.parquet" + ) diff --git a/process_automathtext_arxiv.py b/process_automathtext_arxiv.py new file mode 100644 index 0000000..81324b0 --- /dev/null +++ b/process_automathtext_arxiv.py @@ -0,0 +1,92 @@ +import random +import os +import json +import ast + +from glob import glob +from tqdm import tqdm +import numpy as np +import matplotlib.pyplot as plt +from datasets import load_dataset, Dataset +from huggingface_hub import snapshot_download +from transformers import set_seed + + +# NOTE: Please, store your openai key with "export OPENAI_API_KEY=..." +api_key = os.environ.get("OPENAI_API_KEY") + + +SYSTEM_TEMPLATE = "You are a strong math assistant. You'll be reviewing a math paper and generating questions and answers from the paper." + +# based on Make Your LLM Fully Utilize the Context (https://arxiv.org/pdf/2404.16811) +PROMPT_TEMPLATE = ( + "Generate questions and corresponding answers from the given context. The questions should be highly specific to the " + "mathematical knowledge in the context, not general questions that suits any context.\n\n" + "Rules to follow when generate the questions:\n" + "1. The questions must be fully answerable from mathematical information and knowledge present in given context.\n" + "2. Make sure the questions are clear and unambiguous.\n" + "3. The questions should not focus on the formatting or LaTeX code of the context.\n" + "4. The questions should require the mathematical knowledge and information present in the context.\n\n" + "Rules to follow when generate the answers:\n" + "1. The answers must use the information provided in the context.\n" + "2. The answer should be informative and explain in detail how to arrive at the answer based on the given math content.\n" + "3. Make sure that the explanation includes all the necessary steps and details to arrive at the answer.\n\n" + "Response with {n_qa_pairs} question-answer pairs.\n" + "The question-answer pairs should be in the following format:\n" + "Question 1: {{question_1}}\n" + "Answer 1: {{answer_1}}\n" + "Question 2: {{question_2}}\n" + "Answer 2: {{answer_2}}\n" + "..." + "\n\n" + "### Context ###\n" + "{txt}" +) + + +def get_prompt(txt, n_qa_pairs): + prompt = PROMPT_TEMPLATE.format(txt=txt, n_qa_pairs=n_qa_pairs) + return prompt + + +def get_json_request(txt, n_qa_pairs, gpt_model_name): + prompt = get_prompt(txt, n_qa_pairs) + messages = [ + {"role": "system", "content": SYSTEM_TEMPLATE}, + {"role": "user", "content": prompt}, + ] + return { + "model": gpt_model_name, + "messages": messages, + "temperature": 1.0, + "frequency_penalty": 0.2, + } + + +def remove_too_long(samples): + return [len(text) < 10_000 for text in samples["text"]] + + +if __name__ == "__main__": + set_seed(42) + ds = load_dataset( + "math-ai/AutoMathText", + "arxiv-0.60-to-1.00", + split="train", + trust_remote_code=True, + ) + ds = ds.filter(remove_too_long, batched=True) + print(f"Filtered ds size: {len(ds)}") + + os.makedirs("openai_batches", exist_ok=True) + lines = [] + for sample in tqdm(ds): + code = sample["text"] + if not code: + continue + jsonl = get_json_request(code, n_qa_pairs=2, gpt_model_name="gpt-4o-mini") + lines.append(jsonl) + + with open(f"openai_batches/automathtext_arxiv_qa_pairs.jsonl", "w") as f: + for line in lines: + f.write(json.dumps(line) + "\n") diff --git a/process_automathtext_code_python.py b/process_automathtext_code_python.py new file mode 100644 index 0000000..bcc7745 --- /dev/null +++ b/process_automathtext_code_python.py @@ -0,0 +1,93 @@ +import random +import os +import json +import ast + +from glob import glob +from tqdm import tqdm +import numpy as np +import matplotlib.pyplot as plt +from datasets import load_dataset, Dataset +from huggingface_hub import snapshot_download +from transformers import set_seed + + +# NOTE: Please, store your openai key with "export OPENAI_API_KEY=..." +api_key = os.environ.get("OPENAI_API_KEY") + + +SYSTEM_TEMPLATE = "You are a creative and helpful assistant." + +# based on Make Your LLM Fully Utilize the Context (https://arxiv.org/pdf/2404.16811) +PROMPT_TEMPLATE = ( + "Generate questions and corresponding answers from the given code. The questions should be highly specific to the " + "algorithm implemented in the code, not general questions that suits any code.\n\n" + "Rules to follow when generate the questions:\n" + "1. The questions must be fully answerable from information present in given code. " + "For example, 'What does function `add_two_numbers` do?' or 'What is the purpose of variable `x` in function `add_two_numbers`?'\n" + "2. Make sure the questions are clear and unambiguous.\n\n" + "Rules to follow when generate the answers:\n" + "1. The answers must use the information provided in the code. " + "For example, 'It adds two numbers by using the `+` operator.' or " + "'Variable `x` in function `add_two_numbers` is used to store the result of the addition.'\n\n" + "Response with {n_qa_pairs} question-answer pairs.\n" + "Be creative and think of questions that are not obvious.\n" + "Use simple words and please be clear.\n" + "The question-answer pairs should be in the following format:\n" + "Question 1: {{question_1}}\n" + "Answer 1: {{answer_1}}\n" + "Question 2: {{question_2}}\n" + "Answer 2: {{answer_2}}\n" + "..." + "\n\n" + "### Code ###\n" + "{code}" +) + + +def get_prompt(code, n_qa_pairs): + prompt = PROMPT_TEMPLATE.format(code=code, n_qa_pairs=n_qa_pairs) + return prompt + + +def get_json_request(code, n_qa_pairs, gpt_model_name): + prompt = get_prompt(code, n_qa_pairs) + messages = [ + {"role": "system", "content": SYSTEM_TEMPLATE}, + {"role": "user", "content": prompt}, + ] + return { + "model": gpt_model_name, + "messages": messages, + "temperature": 1.0, + "frequency_penalty": 0.2, + } + + +def remove_too_long(samples): + return [len(text) < 10_000 for text in samples["text"]] + + +if __name__ == "__main__": + set_seed(42) + ds = load_dataset( + "math-ai/AutoMathText", + "code-python-0.80-to-1.00", + split="train", + trust_remote_code=True, + ) + ds = ds.filter(remove_too_long, batched=True) + print(f"Filtered ds size: {len(ds)}") + + os.makedirs("openai_batches", exist_ok=True) + lines = [] + for sample in tqdm(ds): + code = sample["text"] + if not code: + continue + jsonl = get_json_request(code, n_qa_pairs=2, gpt_model_name="gpt-4o-mini") + lines.append(jsonl) + + with open(f"openai_batches/automathtext_code_python_qa_pairs.jsonl", "w") as f: + for line in lines: + f.write(json.dumps(line) + "\n") diff --git a/process_automathtext_web.py b/process_automathtext_web.py new file mode 100644 index 0000000..de5e941 --- /dev/null +++ b/process_automathtext_web.py @@ -0,0 +1,92 @@ +import random +import os +import json +import ast + +from glob import glob +from tqdm import tqdm +import numpy as np +import matplotlib.pyplot as plt +from datasets import load_dataset, Dataset +from huggingface_hub import snapshot_download +from transformers import set_seed + + +# NOTE: Please, store your openai key with "export OPENAI_API_KEY=..." +api_key = os.environ.get("OPENAI_API_KEY") + + +SYSTEM_TEMPLATE = "You are a strong math teaching assistant. You'll be reviewing a web page related to math and generating questions and answers from the page." + +# based on Make Your LLM Fully Utilize the Context (https://arxiv.org/pdf/2404.16811) +PROMPT_TEMPLATE = ( + "Generate questions and corresponding answers from the given context. The questions should be highly specific to the " + "mathematical knowledge in the context, not general questions that suits any context.\n\n" + "Rules to follow when generate the questions:\n" + "1. The questions must be fully answerable from mathematical information present in given context.\n" + "2. Make sure the questions are clear and unambiguous.\n" + "3. The questions should not focus on the formatting or LaTeX code of the context.\n" + "4. The questions should require the mathematical knowledge and information present in the context.\n\n" + "Rules to follow when generate the answers:\n" + "1. The answers must use the information provided in the context.\n" + "2. The answer should be informative and explain in detail how to arrive at the answer based on the given math content.\n" + "3. Make sure that the explanation includes all the necessary steps and details to arrive at the answer.\n\n" + "Response with {n_qa_pairs} question-answer pairs.\n" + "The question-answer pairs should be in the following format:\n" + "Question 1: {{question_1}}\n" + "Answer 1: {{answer_1}}\n" + "Question 2: {{question_2}}\n" + "Answer 2: {{answer_2}}\n" + "..." + "\n\n" + "### Context ###\n" + "{txt}" +) + + +def get_prompt(txt, n_qa_pairs): + prompt = PROMPT_TEMPLATE.format(txt=txt, n_qa_pairs=n_qa_pairs) + return prompt + + +def get_json_request(txt, n_qa_pairs, gpt_model_name): + prompt = get_prompt(txt, n_qa_pairs) + messages = [ + {"role": "system", "content": SYSTEM_TEMPLATE}, + {"role": "user", "content": prompt}, + ] + return { + "model": gpt_model_name, + "messages": messages, + "temperature": 1.0, + "frequency_penalty": 0.2, + } + + +def remove_too_long(samples): + return [len(text) < 10_000 for text in samples["text"]] + + +if __name__ == "__main__": + set_seed(42) + ds = load_dataset( + "math-ai/AutoMathText", + "web-0.80-to-1.00", + split="train", + trust_remote_code=True, + ) + ds = ds.filter(remove_too_long, batched=True) + print(f"Filtered ds size: {len(ds)}") + + os.makedirs("openai_batches", exist_ok=True) + lines = [] + for sample in tqdm(ds): + code = sample["text"] + if not code: + continue + jsonl = get_json_request(code, n_qa_pairs=2, gpt_model_name="gpt-4o-mini") + lines.append(jsonl) + + with open(f"openai_batches/automathtext_web_qa_pairs.jsonl", "w") as f: + for line in lines: + f.write(json.dumps(line) + "\n") diff --git a/process_codeparrot_qa.py b/process_codeparrot_qa.py new file mode 100644 index 0000000..d363284 --- /dev/null +++ b/process_codeparrot_qa.py @@ -0,0 +1,90 @@ +import random +import os +import json +import ast + +from glob import glob +from tqdm import tqdm +import numpy as np +import matplotlib.pyplot as plt +from datasets import load_dataset, Dataset +from huggingface_hub import snapshot_download +from transformers import set_seed + + +# NOTE: Please, store your openai key with "export OPENAI_API_KEY=..." +api_key = os.environ.get("OPENAI_API_KEY") + + +SYSTEM_TEMPLATE = "You are a creative and helpful assistant." + +# based on Make Your LLM Fully Utilize the Context (https://arxiv.org/pdf/2404.16811) +PROMPT_TEMPLATE = ( + "Generate questions and corresponding answers from the given code. The questions should be highly specific to the " + "algorithm implemented in the code, not general questions that suits any code.\n\n" + "Rules to follow when generate the questions:\n" + "1. The questions must be fully answerable from information present in given code. " + "For example, 'What does function `add_two_numbers` do?' or 'What is the purpose of variable `x` in function `add_two_numbers`?'\n" + "2. Make sure the questions are clear and unambiguous.\n\n" + "Rules to follow when generate the answers:\n" + "1. The answers must use the information provided in the code. " + "For example, 'It adds two numbers by using the `+` operator.' or " + "'Variable `x` in function `add_two_numbers` is used to store the result of the addition.'\n\n" + "Response with {n_qa_pairs} question-answer pairs.\n" + "Be creative and think of questions that are not obvious.\n" + "Use simple words and please be clear.\n" + "The question-answer pairs should be in the following format:\n" + "Question 1: {{question_1}}\n" + "Answer 1: {{answer_1}}\n" + "Question 2: {{question_2}}\n" + "Answer 2: {{answer_2}}\n" + "..." + "\n\n" + "### Code ###\n" + "{code}" +) + + +def get_prompt(code, n_qa_pairs): + prompt = PROMPT_TEMPLATE.format(code=code, n_qa_pairs=n_qa_pairs) + return prompt + + +def get_json_request(code, n_qa_pairs, gpt_model_name): + prompt = get_prompt(code, n_qa_pairs) + messages = [ + {"role": "system", "content": SYSTEM_TEMPLATE}, + {"role": "user", "content": prompt}, + ] + return { + "model": gpt_model_name, + "messages": messages, + "temperature": 1.0, + "frequency_penalty": 0.2, + } + + +def remove_too_long(samples): + return [len(text) < 10_000 for text in samples["content"]] + + +if __name__ == "__main__": + set_seed(42) + ds = load_dataset( + "codeparrot/codeparrot-clean-valid", split="train", trust_remote_code=True + ) + ds = ds.filter(remove_too_long, batched=True) + print(f"Filtered ds size: {len(ds)}") + + os.makedirs("openai_batches", exist_ok=True) + lines = [] + for sample in tqdm(ds): + code = sample["content"] + if not code: + continue + jsonl = get_json_request(code, n_qa_pairs=2, gpt_model_name="gpt-4o-mini") + lines.append(jsonl) + + with open(f"openai_batches/code_parrot_qa_pairs.jsonl", "w") as f: + for line in lines: + f.write(json.dumps(line) + "\n") diff --git a/process_fineweb.py b/process_fineweb.py new file mode 100644 index 0000000..5dda98d --- /dev/null +++ b/process_fineweb.py @@ -0,0 +1,37 @@ +import random +from glob import glob + +from datasets import load_dataset, Dataset +from huggingface_hub import snapshot_download +from transformers import set_seed + + +def remove_too_long(samples): + return [len(text) < 10_000 for text in samples["text"]] + + +if __name__ == "__main__": + set_seed(42) + fw_dir = "./data/raw_datasets/fineweb/" + snapshot_download( + "HuggingFaceFW/fineweb", + repo_type="dataset", + local_dir=fw_dir, + allow_patterns="sample/10BT/*", + ) + # https://github.com/huggingface/datasets/issues/7047#issuecomment-2233163406 + num_shards_per_file = 16 + sharded_fw_dir = "./data/raw_datasets/fineweb_sharded/" + output_path_template = f"{sharded_fw_dir}" + "/{index:05d}.parquet" + + for i, f in enumerate(sorted(glob(f"{fw_dir}/sample/10BT/*.parquet"))): + # ~1M rows ~= 2GB mem required per file + ds = Dataset.from_parquet(f) + ds = ds.filter(remove_too_long, batched=True) + print(f"Filtered ds size: {len(ds)}") + ds = ds.shuffle(seed=42 + i) + # take one shard (from 16 shards) per file + idx = random.sample(range(num_shards_per_file), 1)[0] + + shard = ds.shard(index=idx, num_shards=num_shards_per_file, contiguous=False) + shard.to_parquet(output_path_template.format(index=i)) diff --git a/process_openphi_prog.py b/process_openphi_prog.py new file mode 100644 index 0000000..c8bd742 --- /dev/null +++ b/process_openphi_prog.py @@ -0,0 +1,91 @@ +import random +import os +import json +from glob import glob +from tqdm import tqdm +import numpy as np +import matplotlib.pyplot as plt +from datasets import load_dataset, Dataset +from huggingface_hub import snapshot_download +from transformers import set_seed + + +# NOTE: Please, store your openai key with "export OPENAI_API_KEY=..." +api_key = os.environ.get("OPENAI_API_KEY") + + +SYSTEM_TEMPLATE = "You are a creative and helpful assistant." + +# based on Make Your LLM Fully Utilize the Context (https://arxiv.org/pdf/2404.16811) +PROMPT_TEMPLATE = ( + "Generate questions and corresponding answers from the given context. The questions should be highly specific to the " + "information provided in the context, not general questions that suits any context.\n" + "Finally, provide a plausible 1-paragraph continuation of the context. " + "The continuation should reference information in the context in some way.\n\n" + "Rules to follow when generate the questions:\n" + "1. The questions must be fully answerable from information present in given context.\n" + "2. Make sure the questions are clear and unambiguous.\n" + "3. Phrases like 'based on the provided context', 'according to the context', etc, are not allowed to appear in " + "the questions.\n\n" + "Rules to follow when generate the answers:\n" + "1. The answers must use the information provided in the context.\n" + "2. Do not just copy words from the context. Answer the question in your own words.\n\n" + "Response with {n_qa_pairs} question-answer pairs.\n" + "After the question-answer pairs, provide a plausible continuation of the context.\n" + "Use simple words and please be clear.\n" + "The question-answer pairs should be in the following format:\n" + "Question 1: {{question_1}}\n" + "Answer 1: {{answer_1}}\n" + "Question 2: {{question_2}}\n" + "Answer 2: {{answer_2}}\n" + "...\n" + "Continuation: {{continuation}}\n" + "\n\n" + "### Context ###\n" + "{context}" +) + + +def get_prompt(context, n_qa_pairs): + prompt = PROMPT_TEMPLATE.format(context=context, n_qa_pairs=n_qa_pairs) + return prompt + + +def get_json_request(id, text, n_qa_pairs, gpt_model_name): + prompt = get_prompt(text, n_qa_pairs) + messages = [ + {"role": "system", "content": SYSTEM_TEMPLATE}, + {"role": "user", "content": prompt}, + ] + return { + "model": gpt_model_name, + "messages": messages, + "temperature": 1.0, + "frequency_penalty": 0.2, + } + + +def remove_too_long(samples): + return [len(text) < 10_000 for text in samples["markdown"]] + + +if __name__ == "__main__": + set_seed(42) + ds = load_dataset("open-phi/programming_books_llama", split="train") + ds = ds.filter(remove_too_long, batched=True) + print(f"Filtered ds size: {len(ds)}") + + os.makedirs("openai_batches", exist_ok=True) + lines = [] + for i, sample in tqdm(enumerate(ds)): + jsonl = get_json_request( + f"openphi_prog_{i}", + sample["markdown"], + n_qa_pairs=2, + gpt_model_name="gpt-4o-mini", + ) + lines.append(jsonl) + + with open(f"openai_batches/openphi_prog_qa_pairs.jsonl", "w") as f: + for line in lines: + f.write(json.dumps(line) + "\n")