include new data!

.gitignore

@@ -1,5 +1,6 @@
 .cursorrules
 tmp/
+openai_batches/
 models/
 results/
 datasets/

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
 ```

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

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

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)

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 <im_start>{role/name}\n{content}<im_end>\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 <im_start>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),
+        )
+    )

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."},

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

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"

install.sh

@@ -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

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"
+        )

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"
+        )

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")

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")

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")

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")

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))

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")