apunkt/doc-to-lora
1
0
Fork 0
mirror of https://github.com/SakanaAI/doc-to-lora.git synced 2026-07-23 17:01:04 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

rearrange + move to uv (#1)

Browse source
This commit is contained in:
Rujikorn Charakorn 2025-05-27 21:18:15 +09:00 • committed by GitHub
parent 944d844da8
commit 9f986cf938
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
52 changed files with 5114 additions and 5598 deletions
Download patch file Download diff file Expand all files Collapse all files

1
.gitignore vendored
View file

@ -1,3 +1,4 @@
.venv/
.cursorrules
tmp/
openai_batches/

21
.pre-commit-config.yaml
View file

@ -1,17 +1,24 @@
repos:
- repo: https://github.com/asottile/pyupgrade
rev: v2.31.1
rev: v3.20.0
hooks:
- id: pyupgrade
args:
- --py39-plus
- --py310-plus
- repo: https://github.com/astral-sh/ruff-pre-commit
# Ruff version.
rev: v0.11.10
hooks:
# Run the linter.
- id: ruff
types_or: [python, pyi]
args: [--fix]
# Run the formatter.
- id: ruff-format
types_or: [python, pyi]
- repo: https://github.com/PyCQA/isort
rev: 5.12.0
rev: 6.0.1
hooks:
- id: isort
args:
- --profile=black
- repo: https://github.com/python/black
rev: 22.3.0
hooks:
- id: black

2
data/raw_datasets/generate_data.py
View file

@ -1,8 +1,6 @@
import itertools
import json
import os
import random
from typing import Dict, List
def save_jsonl(data: list[dict], filepath: str) -> None:

2
data/raw_datasets/generate_data_big.py
View file

@ -1,8 +1,6 @@
import itertools
import json
import os
import random
from typing import Dict, List
def save_jsonl(data: list[dict], filepath: str) -> None:

39
dl_fineweb_edu.py
View file

@ -1,39 +0,0 @@
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_edu/"
snapshot_download(
"HuggingFaceFW/fineweb-edu",
repo_type="dataset",
local_dir=fw_dir,
allow_patterns="sample/100BT/*",
)
# # 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}" + "/{i:02d}_{idx:05d}.parquet"
# for i, f in enumerate(sorted(glob(f"{fw_dir}/sample/100BT/*.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))
# for idx in range(num_shards_per_file):
# shard = ds.shard(index=idx, num_shards=num_shards_per_file, contiguous=True)
# shard.to_parquet(output_path_template.format(i=i, idx=idx))

34
src/ctx_to_lora/eval.py → eval.py
View file

@ -1,36 +1,36 @@
import re
import string
import yaml
import gc
import json
import os
import logging
import os
import re
import string
from argparse import Namespace
from collections import Counter, defaultdict
from collections.abc import Callable
from dataclasses import fields
from functools import partial
from collections import Counter, defaultdict
from typing import Callable
import torch
import numpy as np
import torch
import yaml
from datasets import disable_caching
from peft import PeftModel
from rouge_score import rouge_scorer
from transformers import (
PreTrainedModel,
EvalPrediction,
GenerationConfig,
Trainer,
PreTrainedModel,
Seq2SeqTrainer,
Seq2SeqTrainingArguments,
EvalPrediction,
Trainer,
set_seed,
)
from transformers.utils import is_liger_kernel_available
from ctx_to_lora.data_definitions import LONGBENCH_E_TASKS, LONGBENCH_TASKS
from ctx_to_lora.data_utils import get_tokenized_dataset
from ctx_to_lora.modeling_utils import ModulatedPretrainedModel
from ctx_to_lora.model_loading import get_tokenizer, get_model
from ctx_to_lora.data.definitions import LONGBENCH_E_TASKS, LONGBENCH_TASKS
from ctx_to_lora.data.processing import get_tokenized_dataset
from ctx_to_lora.model_loading import get_model, get_tokenizer
from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
logger = logging.getLogger()
@ -242,7 +242,9 @@ def decode_test_result(test_dataset, test_result, tokenizer, ctx_tokenizer):
label_toks = sample["labels"][start_idx:]
# labels are padded with -100, so we need to
# replace them with the pad token id
label_toks = np.where(label_toks == -100, tokenizer.pad_token_id, label_toks)
label_toks = np.where(
label_toks == -100, tokenizer.pad_token_id, label_toks
)
label_text = tokenizer.decode(label_toks, skip_special_tokens=True)
d["label"] = label_text.strip()
@ -722,7 +724,7 @@ if __name__ == "__main__":
checkpoint_dir = "/".join(cli_args.checkpoint_path.split("/")[:-1])
run_dir = "/".join(cli_args.checkpoint_path.split("/")[:-2])
cur_it = int(cli_args.checkpoint_path.split("checkpoint-")[1].split("/")[0])
args = Namespace(**yaml.unsafe_load(open(f"{run_dir}/args.yaml", "r")))
args = Namespace(**yaml.unsafe_load(open(f"{run_dir}/args.yaml")))
print(f"checkpoint_path: {cli_args.checkpoint_path}")
print(f"run_dir: {run_dir}")

236
generate_fw_edu_qa_vllm.py
View file

@ -1,236 +0,0 @@
import os
import sys
import random
import re
from glob import glob
import pandas as pd
from datasets import Dataset
from vllm import LLM, SamplingParams
from tqdm import tqdm
from datasets import load_dataset
from openai import OpenAI
SYSTEM_TEMPLATE = (
"You are a creative and helpful assistant.\n"
"You are tasked with generating questions for reading comprehension tests.\n"
"You will be given a context and you need to generate questions and corresponding answers from the given context.\n"
"The questions should be highly specific to the information provided in the context, not general questions that suits any context.\n"
"**DO NOT** halucinate or make up information."
)
# based on Make Your LLM Fully Utilize the Context (https://arxiv.org/pdf/2404.16811)
PROMPT_TEMPLATE = (
"### Instructions ###\n"
"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\n"
"### Context ###\n"
"{context}"
"\n\n\n"
"### Additional Instructions ###\n"
"Rules to follow when generate the questions:\n"
"1. The questions must be specific to the given context and 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"
"4. The questions should not overlap. They should be diverse, covering many aspects of the context.\n"
"5. Do not give away too much information in the questions. For example, ask 'Who is X' instead of 'Who is X that did Y' when Y is clear from the context.\n"
"6. Ignore the text formatting of the context, e.g., bold, italic, underline, etc.\n"
"7. Ignore typos, spacing and grammatical errors 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. Do not just copy words from the context. Answer the question in your own words.\n"
"3. The answers should be detailed and comprehensive. Please include additional specific details from the context when possible.\n\n"
"Response with {n_qa_pairs} question-answer pairs.\n"
"Always use proper grammar and punctuation.\n"
"Try to use different question forms and styles.\n"
"Use simple words and make sure that the answers are clear and comprehensive.\n\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"
"..."
)
def get_prompt(context, n_qa_pairs):
prompt = PROMPT_TEMPLATE.format(context=context, n_qa_pairs=n_qa_pairs)
return prompt
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 length_filter(samples):
return [len(text) < 10_000 for text in samples["text"]]
if __name__ == "__main__":
# api_key = os.environ.get("vllm_api_key")
# vllm_port = os.environ.get("vllm_port")
# vllm_ip = os.environ.get("vllm_ip") # "172.16.0.62"
vllm_model = os.environ.get("vllm_model") # "google/gemma-2-27b-it"
print(f"Using model: {vllm_model}")
# print(f"Using API key: {api_key}")
# print(f"Using VLLM IP: {vllm_ip}")
# print(f"Using VLLM port: {vllm_port}")
# client = OpenAI(
# base_url=f"http://{vllm_ip}:{vllm_port}/v1",
# api_key=api_key,
# )
mm_kwargs = {"image": 0} # if "gemma-3" in vllm_model else {}
llm_kwargs = dict(
model=vllm_model,
dtype="bfloat16",
# tokenizer_mode="mistral",
# config_format="mistral",
# load_format="mistral",
enable_prefix_caching=True,
max_model_len=2**14, # 11264,
limit_mm_per_prompt=mm_kwargs,
# enable_chunked_prefill=True,
)
if vllm_model == "mistralai/Mistral-Small-3.1-24B-Instruct-2503":
llm_kwargs["tokenizer_mode"] = "mistral"
llm_kwargs["config_format"] = "mistral"
llm_kwargs["load_format"] = "mistral"
SYSTEM_TEMPLATE = (
"You are Mistral Small 3.1, a Large Language Model (LLM) created by Mistral AI, a French startup headquartered in Paris.\n"
"You power an AI assistant called Le Chat.\n"
"Your knowledge base was last updated on 2023-10-01.\n\n\n"
) + SYSTEM_TEMPLATE
llm = LLM(**llm_kwargs)
tokenizer = llm.get_tokenizer()
shard_pattern = sys.argv[1]
n_qa_pairs = int(sys.argv[2])
paths = glob(f"./data/raw_datasets/fineweb_edu/sample/100BT/{shard_pattern}.parquet")
for path in paths:
ds = load_dataset(
"parquet",
data_files=path,
split="train",
streaming=True,
)
ds = ds.filter(length_filter, batched=True)
ctxs = [sample["text"] for sample in iter(ds)]
messages = [
[
{"role": "system", "content": SYSTEM_TEMPLATE},
{"role": "user", "content": get_prompt(ctx, n_qa_pairs)},
]
for ctx in ctxs
]
print(f"Generating from {len(messages)} contexts")
completions = llm.chat(
messages,
sampling_params=SamplingParams(
max_tokens=2048,
temperature=(
0.1
if vllm_model == "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
else 0.7
),
# frequency_penalty=0.1,
),
)
samples = []
for ctx, completion in zip(ctxs, completions):
questions, answers = postprocess_qa_pairs(completion.outputs[0].text)
for q, a in zip(questions, answers):
samples.append(
{
"context": ctx,
"prompt": q,
"response": a,
}
)
print(f"Generated {len(samples)} samples")
random.shuffle(samples)
df = pd.DataFrame(samples)
ds = Dataset.from_pandas(df)
val_ds = ds.take(10)
ds = ds.skip(10)
shard_name = path.split("/")[-1].split(".")[0]
ds.to_parquet(f"data/raw_datasets/fw_qa_3/{shard_name}.parquet")
val_ds.to_parquet(f"data/raw_datasets/fw_qa_3/{shard_name}_val.parquet")
print(f"Saved to data/raw_datasets/fw_qa_3/{shard_name}.parquet")
print(f"Saved to data/raw_datasets/fw_qa_3/{shard_name}_val.parquet")
# debug
# ds = load_dataset(
# "parquet",
# data_files=paths[0],
# split="train[:10]",
# )
# ctxs = [sample["text"] for sample in iter(ds)]
# messages = [
# [
# {"role": "system", "content": SYSTEM_TEMPLATE},
# {"role": "user", "content": get_prompt(ctx, n_qa_pairs)},
# ]
# for ctx in ctxs
# ]
#
# print(f"Generating from {len(messages)} contexts")
# completions = llm.chat(
# messages,
# sampling_params=SamplingParams(
# max_tokens=2048,
# temperature=0.1 if vllm_model=="mistralai/Mistral-Small-3.1-24B-Instruct-2503" else 0.7,
# # frequency_penalty=0.1,
# ),
# )
# samples = []
# for ctx, completion in zip(ctxs, completions):
# questions, answers = postprocess_qa_pairs(completion.outputs[0].text)
# for q, a in zip(questions, answers):
# samples.append(
# {
# "context": ctx,
# "prompt": q,
# "response": a,
# }
# )
# for sample in samples:
# print(f"{sample['context']=}")
# print(f"{sample['prompt']=}")
# print(f"{sample['response']=}")
# print()

166
generate_fw_qa.py
View file

@ -1,166 +0,0 @@
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()
res_files = glob("openai_batches/fineweb_qa_pairs_*_res.jsonl")
prompt_files = glob("openai_batches/fineweb_qa_pairs_*[!res].jsonl")
unprocessed_files = set(prompt_files) - set(
[f.replace("_res", "") for f in res_files]
)
for file in unprocessed_files:
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)

231
generate_fw_qa_vllm.py
View file

@ -1,231 +0,0 @@
import os
import sys
import random
import re
from glob import glob
import pandas as pd
from datasets import Dataset
from vllm import LLM, SamplingParams
from tqdm import tqdm
from datasets import load_dataset
from openai import OpenAI
SYSTEM_TEMPLATE = (
"You are a creative and helpful assistant.\n"
"You are tasked with generating questions for reading comprehension tests.\n"
"You will be given a context and you need to generate questions and corresponding answers from the given context.\n"
"The questions should be highly specific to the information provided in the context, not general questions that suits any context.\n"
"**DO NOT** halucinate or make up information."
)
# based on Make Your LLM Fully Utilize the Context (https://arxiv.org/pdf/2404.16811)
PROMPT_TEMPLATE = (
"### Instructions ###\n"
"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\n"
"### Context ###\n"
"{context}"
"\n\n\n"
"### Additional Instructions ###\n"
"Rules to follow when generate the questions:\n"
"1. The questions must be specific to the given context and 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"
"4. The questions should not overlap. They should be diverse, covering many aspects of the context.\n"
"5. Do not give away too much information in the questions. For example, ask 'Who is X' instead of 'Who is X that did Y' when Y is clear from the context.\n"
"6. Ignore the text formatting of the context, e.g., bold, italic, underline, etc.\n"
"7. Ignore typos, spacing and grammatical errors 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. Do not just copy words from the context. Answer the question in your own words.\n"
"3. The answers should be detailed and comprehensive. Please include additional specific details from the context when possible.\n\n"
"Response with {n_qa_pairs} question-answer pairs.\n"
"Always use proper grammar and punctuation.\n"
"Try to use different question forms and styles.\n"
"Use simple words and make sure that the answers are clear and comprehensive.\n\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"
"..."
)
def get_prompt(context, n_qa_pairs):
prompt = PROMPT_TEMPLATE.format(context=context, n_qa_pairs=n_qa_pairs)
return prompt
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 length_filter(samples):
return [len(text) < 10_000 for text in samples["text"]]
if __name__ == "__main__":
# api_key = os.environ.get("vllm_api_key")
# vllm_port = os.environ.get("vllm_port")
# vllm_ip = os.environ.get("vllm_ip") # "172.16.0.62"
vllm_model = os.environ.get("vllm_model") # "google/gemma-2-27b-it"
print(f"Using model: {vllm_model}")
# print(f"Using API key: {api_key}")
# print(f"Using VLLM IP: {vllm_ip}")
# print(f"Using VLLM port: {vllm_port}")
# client = OpenAI(
# base_url=f"http://{vllm_ip}:{vllm_port}/v1",
# api_key=api_key,
# )
mm_kwargs = {"image":0} # if "gemma-3" in vllm_model else {}
llm_kwargs = dict(
model=vllm_model,
# tokenizer_mode="mistral",
# config_format="mistral",
# load_format="mistral",
enable_prefix_caching=True,
max_model_len=12288,#2**14,
limit_mm_per_prompt=mm_kwargs,
# enable_chunked_prefill=True,
)
if vllm_model=="mistralai/Mistral-Small-3.1-24B-Instruct-2503":
llm_kwargs["tokenizer_mode"] = "mistral"
llm_kwargs["config_format"] = "mistral"
llm_kwargs["load_format"] = "mistral"
SYSTEM_TEMPLATE = (
"You are Mistral Small 3.1, a Large Language Model (LLM) created by Mistral AI, a French startup headquartered in Paris.\n"
"You power an AI assistant called Le Chat.\n"
"Your knowledge base was last updated on 2023-10-01.\n\n\n"
) + SYSTEM_TEMPLATE
llm = LLM(**llm_kwargs)
tokenizer = llm.get_tokenizer()
shard_pattern = sys.argv[1]
n_qa_pairs = int(sys.argv[2])
paths = glob(f"./data/raw_datasets/fineweb/sample/100BT/{shard_pattern}.parquet")
# for path in paths:
# ds = load_dataset(
# "parquet",
# data_files=path,
# split="train",
# streaming=True,
# )
# ds = ds.filter(length_filter, batched=True)
#
# ctxs = [sample["text"] for sample in iter(ds)]
# messages = [
# [
# {"role": "system", "content": SYSTEM_TEMPLATE},
# {"role": "user", "content": get_prompt(ctx, n_qa_pairs)},
# ]
# for ctx in ctxs
# ]
#
# print(f"Generating from {len(messages)} contexts")
# completions = llm.chat(
# messages,
# sampling_params=SamplingParams(
# max_tokens=2048,
# temperature=0.1 if vllm_model=="mistralai/Mistral-Small-3.1-24B-Instruct-2503" else 0.7,
# # frequency_penalty=0.1,
# ),
# )
# samples = []
# for ctx, completion in zip(ctxs, completions):
# questions, answers = postprocess_qa_pairs(completion.outputs[0].text)
# for q, a in zip(questions, answers):
# samples.append(
# {
# "context": ctx,
# "prompt": q,
# "response": a,
# }
# )
#
# print(f"Generated {len(samples)} samples")
# random.shuffle(samples)
# df = pd.DataFrame(samples)
# ds = Dataset.from_pandas(df)
# val_ds = ds.take(10)
# ds = ds.skip(10)
#
# shard_name = path.split("/")[-1].split(".")[0]
# ds.to_parquet(f"data/raw_datasets/fw_qa_3/{shard_name}.parquet")
# val_ds.to_parquet(f"data/raw_datasets/fw_qa_3/{shard_name}_val.parquet")
# print(f"Saved to data/raw_datasets/fw_qa_3/{shard_name}.parquet")
# print(f"Saved to data/raw_datasets/fw_qa_3/{shard_name}_val.parquet")
# debug
ds = load_dataset(
"parquet",
data_files=paths[0],
split="train[:10]",
)
ctxs = [sample["text"] for sample in iter(ds)]
messages = [
[
{"role": "system", "content": SYSTEM_TEMPLATE},
{"role": "user", "content": get_prompt(ctx, n_qa_pairs)},
]
for ctx in ctxs
]
print(f"Generating from {len(messages)} contexts")
completions = llm.chat(
messages,
sampling_params=SamplingParams(
max_tokens=2048,
temperature=0.1 if vllm_model=="mistralai/Mistral-Small-3.1-24B-Instruct-2503" else 0.7,
# frequency_penalty=0.1,
),
)
samples = []
for ctx, completion in zip(ctxs, completions):
questions, answers = postprocess_qa_pairs(completion.outputs[0].text)
for q, a in zip(questions, answers):
samples.append(
{
"context": ctx,
"prompt": q,
"response": a,
}
)
for sample in samples:
print(f"{sample['context']=}")
print(f"{sample['prompt']=}")
print(f"{sample['response']=}")
print()

44
generate_negative_natural_questions.py
View file

@ -1,44 +0,0 @@
import os
from tqdm import tqdm
from datasets import load_dataset, Dataset
if __name__ == "__main__":
ds = load_dataset("google-research-datasets/natural_questions")
ds = ds.shuffle(seed=42)
for split in ds:
out = []
for i, sample in enumerate(tqdm(ds[split])):
ctx = " ".join(
[
token
for is_html, token in zip(
sample["document"]["tokens"]["is_html"],
sample["document"]["tokens"]["token"],
)
if not is_html
]
).strip()
if not ctx:
continue
answers = []
for answer in sample["annotations"]["short_answers"]:
answers += answer["text"]
if not answers:
continue
answer = answers[0].strip()
if not answer:
continue
prompt = sample["question"]["text"].capitalize().strip() + " ?"
out.append(dict(prompt=prompt, context=ctx, answer=answer.capitalize()))
# if i >= 10:
# break
shifted_ctxs = [d["context"] for d in out]
shifted_ctxs = shifted_ctxs[1:] + shifted_ctxs[:1]
for d, shifted_ctx in zip(out, shifted_ctxs):
d["context"] = shifted_ctx
new_ds = Dataset.from_list(out)
print(new_ds)
save_dir = "data/raw_datasets/negative_natural_questions"
os.makedirs(save_dir, exist_ok=True)
new_ds.to_json(f"{save_dir}/{split}.jsonl")

22
generate_one_big_fw_qa_req.py
View file

@ -1,22 +0,0 @@
import json
from glob import glob
if __name__ == "__main__":
res_files = glob("openai_batches/fineweb_qa_pairs_*_res.jsonl")
prompt_files = glob("openai_batches/fineweb_qa_pairs_*[!res].jsonl")
unprocessed_files = set(prompt_files) - set(
[f.replace("_res", "") for f in res_files]
)
print(f"Found {len(unprocessed_files)} unprocessed files:\n{unprocessed_files}")
# Concatenate "body" from each line from unprocessed files
with open("openai_batches/fineweb_qa_pairs_large.jsonl", "w") as outfile:
for fname in unprocessed_files:
with open(fname) as infile:
for line in infile:
# Load the JSON object from the line
obj = json.loads(line)
# Extract the "body" and write it as a new line
if "body" in obj:
outfile.write(json.dumps(obj["body"]) + "\n")
print(f"Concatenated 'body' from {len(unprocessed_files)} files")

495
generate_qa_parallel.py
View file

@ -1,495 +0,0 @@
# 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"))
# tier 4 limit
parser.add_argument("--max_requests_per_minute", type=int, default=10_000 * 0.5)
parser.add_argument("--max_tokens_per_minute", type=int, default=10_000_000 * 0.5)
#
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),
)
)

516
generate_qa_parallel_vllm.py
View file

@ -1,516 +0,0 @@
# 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"))
# tier 4 limit
parser.add_argument("--max_requests_per_minute", type=int, default=10_000 * 0.5)
parser.add_argument("--max_tokens_per_minute", type=int, default=10_000_000 * 0.5)
#
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),
)
)
"""
from openai import OpenAI
import os
api_key = os.environ.get("vllm_api_key")
vllm_port = os.environ.get("vllm_port")
vllm_ip = "172.16.0.62"
client = OpenAI(
base_url=f"http://{vllm_ip}:{vllm_port}/v1",
api_key=api_key,
)
completion = client.chat.completions.create(
model="google/gemma-2-27b-it",
messages=[
{"role": "user", "content": "Hello!"}
]
)
print(completion.choices[0].message)
"""

8
icae_v2/fine_tuned_inference.py
View file

@ -1,5 +1,4 @@
import json
import sys
import torch
from modeling_icae_multi_span import (
@ -11,7 +10,7 @@ from modeling_icae_multi_span import (
from peft import LoraConfig
from safetensors.torch import load_file
from tqdm import tqdm
from transformers import AutoModelForCausalLM, HfArgumentParser
from transformers import HfArgumentParser
# Set the computation device
device = "cuda"
@ -53,7 +52,6 @@ model.eval()
with torch.no_grad():
with open("ft_inference.out", "w") as f:
for line in tqdm(lines):
# Tokenize input text
data = json.loads(line)
@ -100,7 +98,9 @@ with torch.no_grad():
# Generate text output
for i in range(max_out_length):
with model.icae.disable_adapter(): # no independent decoder; use self.icae
with (
model.icae.disable_adapter()
): # no independent decoder; use self.icae
out = model.icae(
inputs_embeds=output,
past_key_values=past_key_values,

7
icae_v2/instruction_finetune.py
View file

@ -7,7 +7,6 @@ from modeling_icae_multi_span import (
)
from peft import LoraConfig
from training_utils import (
DataCollatorForDynamicPadding,
instruct_ft_tokenize_function,
train_model,
)
@ -72,9 +71,9 @@ def main():
)
# check model_args.mem_size and min_tokens_for_lm
assert (
training_args.fixed_mem_size & (training_args.fixed_mem_size - 1)
) == 0, "training_args.fixed_mem_size must be a power of 2"
assert (training_args.fixed_mem_size & (training_args.fixed_mem_size - 1)) == 0, (
"training_args.fixed_mem_size must be a power of 2"
)
memory_size = training_args.fixed_mem_size

32
icae_v2/modeling_icae_multi_span.py
View file

@ -1,18 +1,14 @@
# ICAE that supports multi span concat
import math
import os
import random
from dataclasses import dataclass, field
from typing import Optional
import torch
import torch.nn as nn
import transformers
from peft import get_peft_model
from safetensors.torch import load_file
from torch.nn.functional import gelu
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
from transformers import AutoModelForCausalLM, AutoTokenizer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@ -32,16 +28,20 @@ class ModelArguments:
@dataclass
class DataArguments:
data_path: str = field(default=None, metadata={"help": "Path to the training data."})
data_path: str = field(
default=None, metadata={"help": "Path to the training data."}
)
debug_data: bool = field(
default=False,
metadata={"help": "Enable debug dataset to quickly verify the training process"},
metadata={
"help": "Enable debug dataset to quickly verify the training process"
},
)
@dataclass
class TrainingArguments(transformers.TrainingArguments):
cache_dir: Optional[str] = field(default=None)
cache_dir: str | None = field(default=None)
optim: str = field(default="adamw_torch")
model_max_length: int = field(
default=28000,
@ -77,7 +77,9 @@ class TrainingArguments(transformers.TrainingArguments):
)
restore_from: str = field(
default="",
metadata={"help": "The checkpoint that should be restored from for fine-tuning"},
metadata={
"help": "The checkpoint that should be restored from for fine-tuning"
},
)
@ -109,7 +111,9 @@ class ICAE(torch.nn.Module):
self.model_name = model_args.model_name_or_path
self.icae = AutoModelForCausalLM.from_pretrained(
self.model_name,
torch_dtype=torch.float16 if training_args.bf16 is False else torch.bfloat16,
torch_dtype=torch.float16
if training_args.bf16 is False
else torch.bfloat16,
use_flash_attention_2=True,
resume_download=True,
)
@ -162,9 +166,7 @@ class ICAE(torch.nn.Module):
self.vocab_size + self.mem_size,
dtype=torch.long,
device=device,
).unsqueeze(
0
) # mem tokens
).unsqueeze(0) # mem tokens
if self.training:
self.init()
@ -201,7 +203,7 @@ class ICAE(torch.nn.Module):
self,
input_ids: torch.LongTensor = None,
prompt_answer_ids: torch.LongTensor = None,
labels: Optional[torch.LongTensor] = None,
labels: torch.LongTensor | None = None,
):
# encoder part
batch_size = input_ids.size(0)
@ -218,7 +220,6 @@ class ICAE(torch.nn.Module):
)
for segment_idx in range(num_segments):
start_idx = segment_idx * segment_length
end_idx = min((segment_idx + 1) * segment_length, total_length)
segment_input_ids = input_ids[:, start_idx:end_idx]
@ -292,7 +293,6 @@ class ICAE(torch.nn.Module):
def _compress(
self, input_ids: torch.LongTensor = None
): # for inference; compress a fixed length of input into memory slots
batch_size = input_ids.size(0)
total_length = input_ids.size(1)
num_segments = self.compute_num_segments(total_length)

12
icae_v2/pretrain.py
View file

@ -36,12 +36,12 @@ def main():
)
# check model_args.mem_size and min_tokens_for_lm
assert (
training_args.fixed_mem_size & (training_args.fixed_mem_size - 1)
) == 0, "training_args.fixed_mem_size must be a power of 2"
assert (
training_args.leave_tokens_for_lm <= training_args.min_tokens_for_lm
), "leave_tokens_for_lm should be fewer than min_tokens_for_lm"
assert (training_args.fixed_mem_size & (training_args.fixed_mem_size - 1)) == 0, (
"training_args.fixed_mem_size must be a power of 2"
)
assert training_args.leave_tokens_for_lm <= training_args.min_tokens_for_lm, (
"leave_tokens_for_lm should be fewer than min_tokens_for_lm"
)
memory_size = training_args.fixed_mem_size

6
icae_v2/pretrained_inference.py
View file

@ -1,5 +1,4 @@
import json
import sys
import torch
from modeling_icae_multi_span import (
@ -49,7 +48,6 @@ lines = [
model.eval()
with torch.no_grad():
with open("tmp.out", "w") as f:
for line in tqdm(lines):
# Tokenize input text
tokenized_text = model.tokenizer(
@ -80,7 +78,9 @@ with torch.no_grad():
# Generate text output
for i in range(512):
with model.icae.disable_adapter(): # no independent decoder; use self.icae
with (
model.icae.disable_adapter()
): # no independent decoder; use self.icae
out = model.icae(
inputs_embeds=output,
past_key_values=past_key_values,

11
icae_v2/training_utils.py
View file

@ -1,9 +1,8 @@
import math
import os
import random
import torch
from transformers import Seq2SeqTrainer, Trainer
from transformers import Trainer
from transformers.trainer_utils import get_last_checkpoint
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@ -17,7 +16,6 @@ def train_model(
data_collator=None,
compute_metrics=None,
):
last_checkpoint = None
if (
os.path.isdir(training_args.output_dir)
@ -85,7 +83,6 @@ def train_model(
def text_extraction(input_ids, length, lm_ratio=0.0):
input_len = len(input_ids)
assert input_len >= 1, f"Error: invalid input length ({input_len})"
@ -121,9 +118,10 @@ def pretrain_tokenize_function(examples, model, mem, lm_ratio=0.0):
max_len = model.training_args.model_max_length # heuristic
for idx in range(len(text_output["input_ids"])):
ae = True
a, b = text_extraction(text_output["input_ids"][idx], max_len, lm_ratio=lm_ratio)
a, b = text_extraction(
text_output["input_ids"][idx], max_len, lm_ratio=lm_ratio
)
length_a = len(a)
num_segments = model.compute_num_segments(length_a)
total_mem_length = num_segments * model.mem_size
@ -192,7 +190,6 @@ def instruct_ft_tokenize_function(examples, model, mem):
max_len = model.training_args.model_max_length # heuristic
for idx in range(len(text_output["input_ids"])):
length = len(text_output["input_ids"][idx])
num_segments = model.compute_num_segments(length)
total_mem_length = num_segments * model.mem_size

25
install.sh
View file

@ -1,9 +1,18 @@
pip install -e .
pip install flash-attn==2.6.3 --no-build-isolation
# huggingface-cli login
# uv pip install -e .
# uv pip install flash-attn==2.6.3 --no-build-isolation
# # huggingface-cli login
# lm-evaluation-harness
git clone https://github.com/EleutherAI/lm-evaluation-harness.git
cd lm-evaluation-harness
git checkout f724be699e8adf7ca8004ea0e519dfac83a06f18
pip install -e .
# # lm-evaluation-harness
# # git clone https://github.com/EleutherAI/lm-evaluation-harness.git
# # cd lm-evaluation-harness
# # git checkout f724be699e8adf7ca8004ea0e519dfac83a06f18
# # pip install -e .
uv venv
uv pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cu126
uv sync
# wget https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.3/flash_attn-2.7.3+cu12torch2.6cxx11abiTRUE-cp310-cp310-linux_x86_64.whl
uv pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.3/flash_attn-2.7.3+cu12torch2.6cxx11abiTRUE-cp310-cp310-linux_x86_64.whl
# uv pip install flash-attn==2.6.3 --no-build-isolation
# dev
pre-commit install

365
intx_sft.py
View file

@ -1,68 +1,57 @@
import logging
import os
import random
import string
import time
from multiprocess import set_start_method
from math import ceil
from collections import defaultdict
from copy import copy, deepcopy
from copy import deepcopy
from functools import partial
from importlib.resources import read_binary
from typing import Callable, Optional
from math import ceil
import numpy as np
import torch
import wandb
import yaml
from ctx_to_lora.data_utils import (
DS_LEN,
convert_ctx_prompt_response_to_messages,
get_preprocessing_fn,
get_sft_prompt_formatting_fn,
get_tokenized_dataset,
tokenize_chat_messages,
tokenize_ctx_text,
)
from datasets import (
concatenate_datasets,
interleave_datasets,
disable_caching,
load_dataset,
IterableDataset,
interleave_datasets,
)
from peft import PeftModel
from transformers import (
AutoConfig,
set_seed,
)
from transformers.utils import is_liger_kernel_available
import wandb
from ctx_to_lora.configs import (
AggregatorArguments,
ArgumentParser,
CtxEncoderArguments,
CtxTrainingArguments,
DataArguments,
ExperimentSetup,
HypernetArguments,
LoRAArguments,
ModelArguments,
TrainingArguments,
)
from ctx_to_lora.data.collator import train_collator, train_packed_collator
from ctx_to_lora.data.definitions import DS_LEN
from ctx_to_lora.data.processing import get_tokenized_dataset
from ctx_to_lora.metrics import (
Evaluator,
compute_metrics,
compute_per_token_acc,
compute_perplexity,
compute_prefix_matching,
)
from ctx_to_lora.model_loading import (
get_lora_config,
get_model_and_tokenizer,
get_tokenizer,
)
from ctx_to_lora.modeling_utils import (
EarlyExit,
HyperLoRA,
from ctx_to_lora.modeling.hypernet import (
ModulatedPretrainedModel,
get_hypernet_config,
)
from rouge_score import rouge_scorer
from ctx_to_lora.training_utils import TRAINING_TASK, train_model
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
AutoConfig,
PretrainedConfig,
DataCollatorForSeq2Seq,
EvalPrediction,
HfArgumentParser,
set_seed,
)
from accelerate import Accelerator, DeepSpeedPlugin
from peft import PeftModel
from transformers.utils import is_liger_kernel_available
from transformers.data import DataCollatorWithFlattening
from torch.utils.data import DataLoader
from ctx_to_lora.trainer import train_model
from ctx_to_lora.utils import (
extract_cli_args,
get_base_model,
get_run_name,
log_num_train_params,
save_yaml,
@ -70,150 +59,11 @@ from ctx_to_lora.utils import (
validate_args,
)
from ctx_to_lora.configs import (
ArgumentParser,
CtxTrainingArguments,
TrainingArguments,
DataArguments,
ExperimentSetup,
LoRAArguments,
ModelArguments,
HypernetArguments,
AggregatorArguments,
CtxEncoderArguments,
)
logger = logging.getLogger()
LOCAL_RANK = int(os.getenv("LOCAL_RANK", "0"))
def compute_per_token_acc(shift_logits, shift_labels, valid_masks):
indices = torch.where(valid_masks)
# acc = (shift_logits.argmax(-1) == shift_labels)[indices].float().mean().item()
# return {"per_token_acc": acc}
acc = (shift_logits.argmax(-1) == shift_labels)[indices].float()
return {
"per_token_accs": acc.flatten().tolist(),
"n_per_token_accs": valid_masks.sum().item(),
}
def compute_prefix_matching(shift_logits, shift_labels, valid_masks):
lengths = valid_masks.sum(dim=1)
is_wrong = (shift_logits.argmax(-1) != shift_labels) * valid_masks
is_correct = (shift_logits.argmax(-1) == shift_labels) * valid_masks
# NOTE: not reliable for multi-turn conversations
# ie, all tokens in the following user's turn will be correct
# still monotonically correlate with perf though
wrong_pos = torch.argmax(is_wrong, dim=1) - torch.argmax(valid_masks, dim=1)
perf = wrong_pos / lengths
# if all tokens are correct, set to 1
perf = torch.where(is_correct.sum(dim=1) == lengths, 1, perf)
# return {"prefix_matching": perf.mean().item()}
return {
"prefix_matchings": perf.tolist(),
"n_prefix_matchings": valid_masks.shape[0],
}
@torch.no_grad()
def compute_perplexity(shift_logits, shift_labels, valid_masks):
loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
loss = loss_fct(shift_logits.transpose(1, 2), shift_labels)
loss = (loss * valid_masks).sum(dim=1) / valid_masks.sum(dim=1)
# perplexity = torch.exp(loss).mean().item()
# return {"perplexity": perplexity}
preplexities = torch.exp(loss)
return {
"perplexities": preplexities.tolist(),
"n_perplexities": valid_masks.shape[0],
}
class Evaluator:
def __init__(self, metric_fns: list[Callable]):
self.metric_fns = metric_fns
self.reset()
def reset(self):
self.accum_metrics = defaultdict(list)
self.count = defaultdict(list)
def update(self, shift_logits, shift_labels, valid_masks):
for metric_fn in self.metric_fns:
metric = metric_fn(shift_logits, shift_labels, valid_masks)
for k, v in metric.items():
if k.startswith("n_"):
self.count[k[2:]].append(v)
else:
self.accum_metrics[k] += v
def compute(self):
# Get result across entire eval set
result = {
k: np.sum(v) / np.sum(self.count[k]) for k, v in self.accum_metrics.items()
}
# Reset batch statistics
self.reset()
return result
@torch.no_grad()
def compute_metrics(
eval_pred: EvalPrediction,
compute_result: bool,
evaluator: Evaluator,
) -> Optional[dict]:
logits, labels = eval_pred.predictions, eval_pred.label_ids
shift_logits = logits[..., :-1, :]
shift_labels = labels[..., 1:]
valid_masks = torch.where(shift_labels != -100, 1, 0)
evaluator.update(shift_logits, shift_labels, valid_masks)
if compute_result:
return evaluator.compute()
# def compute_metrics(eval_pred: EvalPrediction) -> dict:
# """
# Custom metrics function for the trainer
# Args:
# eval_pred: tuple of predictions and labels
# Returns:
# dictionary containing metric names (str) and values (Any)
# """
# # compute per token accuracy
# # logits, labels = eval_pred.predictions, eval_pred.label_ids
# # shift_logits = logits[..., :-1, :]
# pred_ids, labels = eval_pred.predictions, eval_pred.label_ids
# shift_pred_ids = pred_ids[..., :-1]
# shift_labels = labels[..., 1:]
# valid_masks = np.where(shift_labels != -100, 1, 0)
# per_token_acc = compute_per_token_acc(shift_pred_ids, shift_labels, valid_masks)
# prefix_matching = compute_prefix_matching(shift_pred_ids, shift_labels, valid_masks)
# # entropy = compute_entropy(shift_logits, shift_labels, valid_masks)
# return dict(
# **per_token_acc,
# **prefix_matching,
# # **entropy,
# num_valid_tokens=valid_masks.sum(),
# num_samples=valid_masks.shape[0],
# )
# # https://discuss.huggingface.co/t/cuda-out-of-memory-when-using-trainer-with-compute-metrics/2941/13
# def preprocess_logits_for_metrics(logits, labels):
# """
# Original Trainer may have a memory leak.
# This is a workaround to avoid storing too many tensors that are not needed.
# """
# pred_ids = torch.argmax(logits, dim=-1)
# return pred_ids
def main():
############ Argument parsing
parser = ArgumentParser(
@ -302,7 +152,6 @@ def main():
save_yaml(args, f"{output_dir}/args.yaml")
############ Model setup
if not ctx_args.from_pretrained_checkpoint:
model_name = model_args.model_name_or_path
model, tokenizer = get_model_and_tokenizer(
@ -364,6 +213,10 @@ def main():
raise ValueError("ctx_encoder contains trainable parameters")
if len([p for p in model.base_model.parameters() if p.requires_grad]):
raise ValueError("base model contains trainable parameters")
# we can't compile the base model bc we will be
# adding/removing forward hooks during training
model.hypernet = torch.compile(model.hypernet)
model.ctx_encoder = torch.compile(model.ctx_encoder)
else:
# activate LoRA
base_model_config = AutoConfig.from_pretrained(
@ -372,29 +225,27 @@ def main():
base_model_config.save_pretrained(output_dir)
logger.info("Using LoRA")
model.set_adapter("default")
model = torch.compile(model)
model.train()
logger.debug(model)
log_num_train_params(model)
# model = accelerator.prepare(model)
############ Dataset setup
logger.info("Loading dataset...")
add_ctx_to_chat = not isinstance(model, ModulatedPretrainedModel)
# tokenizer_kwargs = {"max_length": ctx_args.max_base_len}
# ctx_tokenizer_kwargs = {"max_length": ctx_args.max_ctx_len} # not used for now
tokenizer_kwargs = {"max_length": ctx_args.max_base_len} # not used
ctx_tokenizer_kwargs = {"max_length": ctx_args.max_ctx_len} # not used for now
_get_tokenized_dataset = partial(
get_tokenized_dataset,
base_model_max_len=model.base_model.config.max_position_embeddings,
tokenizer=tokenizer,
tokenizer_kwargs={},
tokenizer_kwargs=tokenizer_kwargs,
ctx_model_max_len=model.ctx_encoder.config.max_position_embeddings,
ctx_tokenizer=ctx_tokenizer,
ctx_tokenizer_kwargs={},
ctx_tokenizer_kwargs=ctx_tokenizer_kwargs,
add_ctx_to_chat=add_ctx_to_chat,
add_repeat_prompt=ctx_args.add_repeat_prompt,
add_negative_prompt=ctx_args.add_negative_prompt,
@ -411,7 +262,6 @@ def main():
continue
streaming = (split == "train") and data_args.streaming
for ds_name in ds_names:
# with accelerator.main_process_first():
ds = _get_tokenized_dataset(ds_name, split, streaming=streaming)
tokenized_ds[split][os.path.basename(ds_name)] = ds
@ -431,86 +281,50 @@ def main():
val_indices = np.random.permutation(len(ds))[:n_val_samples]
val_ds[ds_name] = val_ds[ds_name].select(val_indices)
max_steps = ceil(
sum([DS_LEN[ds] for ds in train_ds])
* training_args.num_train_epochs
/ training_args.per_device_train_batch_size
/ training_args.gradient_accumulation_steps
/ training_args.world_size
)
training_args.max_steps = max_steps
if data_args.streaming:
max_steps = ceil(
sum(DS_LEN[ds] for ds in train_ds)
* training_args.num_train_epochs
/ training_args.per_device_train_batch_size
/ training_args.gradient_accumulation_steps
/ training_args.world_size
)
training_args.max_steps = max_steps
# interleaving streaming datasets
# simplify the probs for smaller datasets
# slightly upsample those datasets
probs = [
0.01 if "fw_qa" not in ds_name else 1 + 0.01 - 0.01 * len(train_ds)
for ds_name in train_ds
]
train_ds = interleave_datasets(
list(train_ds.values()),
probabilities=probs,
stopping_strategy="all_exhausted",
seed=training_args.seed,
)
# interleaving streaming datasets
# simplify the probs for smaller datasets
# slightly upsample those datasets
probs = [
0.01 if "fw_qa" not in ds_name else 1 + 0.01 - 0.01 * len(train_ds)
for ds_name in train_ds
]
train_ds = interleave_datasets(
list(train_ds.values()),
probabilities=probs,
stopping_strategy="all_exhausted",
seed=training_args.seed,
)
else:
train_ds_len = [len(ds) for ds in train_ds.values()]
total_len = sum(train_ds_len)
max_steps = ceil(
total_len
* training_args.num_train_epochs
/ training_args.per_device_train_batch_size
/ training_args.gradient_accumulation_steps
/ training_args.world_size
)
training_args.max_steps = max_steps
train_ds = interleave_datasets(
list(train_ds.values()),
probabilities=[l / total_len for l in train_ds_len],
seed=training_args.seed,
)
logger.info(f"train_ds: {train_ds}")
logger.info(f"val_ds: {val_ds}")
flattener = DataCollatorWithFlattening()
def train_packed_collator(inp_list):
# no padding
packed_inputs = flattener(inp_list, return_tensors="pt")
if "ctx_ids" in inp_list[0]:
ctx_ids = [{"input_ids": example["ctx_ids"]} for example in inp_list]
packed_ctx = flattener(ctx_ids, return_tensors="pt")
packed_inputs["ctx_ids"] = packed_ctx["input_ids"]
packed_inputs["ctx_position_ids"] = packed_ctx["position_ids"]
return packed_inputs
def train_collator(inp_list, tokenizer):
# input is a list of tokenized sequences
padding_kwargs = dict(
padding=True,
padding_side="right",
pad_to_multiple_of=8,
return_tensors="pt",
)
ctx_ids = None
if "ctx_ids" in inp_list[0]:
# have to be manual since it has [ctx_len, features] shape
# pad to the longest ctx_len in the batch
# which can have a different length from the input_ids, attn_mask, labels
ctx_ids = [example.pop("ctx_ids") for example in inp_list]
ctx_ids = torch.nn.utils.rnn.pad_sequence(
ctx_ids,
batch_first=True,
padding_value=0,
)
# exotic keys won't be padded, so we need to pad them as well
ctx_attn_mask = [example.pop("ctx_attn_mask") for example in inp_list]
ctx_attn_mask = torch.nn.utils.rnn.pad_sequence(
ctx_attn_mask,
batch_first=True,
padding_value=0,
)
labels = [x.pop("labels") for x in inp_list]
padded_seq = tokenizer.pad(inp_list, **padding_kwargs)
# hacky explicit padding since the labels are not padded by default
labels = tokenizer.pad({"input_ids": labels}, **padding_kwargs)["input_ids"]
labels = torch.where(padded_seq["attention_mask"] == 0, -100, labels)
out = {**padded_seq, "labels": labels}
if ctx_ids is not None:
out["ctx_ids"] = ctx_ids
out["ctx_attn_mask"] = ctx_attn_mask
return out
train_collator = (
collator = (
train_packed_collator
if ctx_args.use_sequence_packing
else partial(train_collator, tokenizer=tokenizer)
@ -553,20 +367,16 @@ def main():
train_model(
model,
# tokenizer,
training_args,
train_ds,
val_ds,
train_collator,
# partial(generation_collator, tokenizer=tokenizer),
collator,
compute_metrics=partial(
compute_metrics,
evaluator=Evaluator(
[compute_per_token_acc, compute_prefix_matching, compute_perplexity]
),
),
# max_new_tokens=ctx_args.max_new_tokens,
# gen_per_device_eval_batch_size=ctx_args.gen_per_device_eval_batch_size,
)
logger.info(f"Training run finished and saved to {output_dir}")
@ -574,6 +384,7 @@ def main():
if __name__ == "__main__":
os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true"
os.environ["TOKENIZERS_PARALLELISM"] = "true"
os.environ["WANDB_DIR"] = ".wandb/"
os.environ["WANDB_PROJECT"] = "ctx_to_lora"
os.environ["WANDB_WATCH"] = "" # "all"
os.environ["WANDB_CONSOLE"] = "off"
@ -581,9 +392,9 @@ if __name__ == "__main__":
# os.environ["KMP_AFFINITY"] = "disabled" # fixing iterable dataset stuck
os.environ["OMP_NUM_THREADS"] = "23"
# os.environ["HF_DATASETS_IN_MEMORY_MAX_SIZE"] = "137438953472" # 128 TB
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch._dynamo.config.capture_scalar_outputs = True
if os.getenv("DEBUG", False):
disable_caching()
# randomly sleep to avoid run_name collision
# time.sleep(random.random() * 13)
# set_start_method("spawn", force=True)
main()

216
post_process_fw_qa.py
View file

@ -1,216 +0,0 @@
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"
)

212
post_process_parallel_qa.py
View file

@ -1,212 +0,0 @@
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"
)
# val_ds = ds.take(500)
# ds = ds.skip(500)
# ds.to_parquet(f"data/raw_datasets/fw_qa_large/{fname}.parquet")
# val_ds.to_parquet(f"data/raw_datasets/fw_qa_large/{fname}_val.parquet")
# print(f"Saved {fname} to data/raw_datasets/fw_qa_large/{fname}.parquet")
# print(f"Saved {fname}_val to data/raw_datasets/fw_qa_large/{fname}_val.parquet")

92
process_automathtext_arxiv.py
View file

@ -1,92 +0,0 @@
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")

93
process_automathtext_code_python.py
View file

@ -1,93 +0,0 @@
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")

92
process_automathtext_web.py
View file

@ -1,92 +0,0 @@
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")

90
process_codeparrot_qa.py
View file

@ -1,90 +0,0 @@
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")

39
process_fineweb.py
View file

@ -1,39 +0,0 @@
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/100BT/*",
)
# # 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}" + "/{i:02d}_{idx:05d}.parquet"
# for i, f in enumerate(sorted(glob(f"{fw_dir}/sample/100BT/*.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))
# for idx in range(num_shards_per_file):
# shard = ds.shard(index=idx, num_shards=num_shards_per_file, contiguous=True)
# shard.to_parquet(output_path_template.format(i=i, idx=idx))

91
process_openphi_prog.py
View file

@ -1,91 +0,0 @@
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")

30
pyproject.toml
View file

@ -10,8 +10,8 @@ dependencies = [
"deepspeed==0.16.7",
"accelerate==1.6.0",
"datasets==3.6.0",
"setuptools",
"peft",
"bitsandbytes",
"jupyter",
"matplotlib",
"hf_transfer",
@ -35,15 +35,22 @@ dependencies = [
requires = ["setuptools>=61.0"]
build-backend = "setuptools.build_meta"
[tool.isort]
profile = "black"
[tool.black]
line-length = 89
[tool.pyright]
exclude = [
"**/node_modules",
"**/__pycache__",
"**/.*",
".venv",
".github",
".vscode",
"chat_templates",
"eval_results",
"configs",
"EditingLlama",
"icae_v2",
"lm-evaluation-harness",
"LongBench",
"scripts",
"train_outputs",
"data",
"generated_tasks",
@ -55,4 +62,9 @@ exclude = [
typeCheckingMode = "off"
[tool.ruff]
line-length = 89
line-length = 88
select = ["F401"] # remove unused imports
ignore = ["E", "F"]
[tool.isort]
profile = "black"

13
requirements.txt
View file

@ -1,13 +0,0 @@
einops
jaxtyping
liger-kernel
rouge-score
transformers==4.51.1
# vllm==0.8.3
deepspeed==0.15.4 # curretly 0.16.2
accelerate==1.2.1
# visualization
tensorboard
flask
pandas
plotly

9
run.sh
View file

@ -1,9 +0,0 @@
job_id=$(sbatch --parsable batch.sh "$@")
while ! scontrol show job "$job_id" | grep -q "JobState=RUNNING"; do
echo "Waiting for job $job_id to start..."
sleep 5
done
echo "Job $job_id started"
sattach $job_id.0

19
src/ctx_to_lora/run_lm_eval.py → run_lm_eval.py
View file

@ -1,15 +1,14 @@
import sys
import os
import json
import torch
import lm_eval
from lm_eval.models.huggingface import HFLM
from lm_eval.models.vllm_causallms import VLLM
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel, PeftConfig
import argparse
import json
import os
from ctx_to_lora.modeling_utils import (
import lm_eval
import torch
from lm_eval.models.huggingface import HFLM
from peft import PeftConfig, PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer
from ctx_to_lora.modeling.hypernet import (
ModulatedModelWithSharedInput,
ModulatedPretrainedModel,
)

2
setup.py
View file

@ -1,7 +1,7 @@
# read the contents of the README file
from pathlib import Path
from setuptools import find_packages, setup
from setuptools import setup
this_directory = Path(__file__).parent
long_description = (this_directory / "README.md").read_text()

51
src/ctx_to_lora/configs.py
View file

@ -2,12 +2,16 @@ import dataclasses
import os
import sys
from dataclasses import dataclass, field
from enum import Enum, auto
from typing import Any, Dict, List, Literal, NewType, Optional, Tuple
from enum import Enum
from typing import Any, Literal, NewType
import torch
import yaml
from transformers import MODEL_FOR_CAUSAL_LM_MAPPING, HfArgumentParser, TrainingArguments
from transformers import (
MODEL_FOR_CAUSAL_LM_MAPPING,
HfArgumentParser,
TrainingArguments,
)
MODEL_CONFIG_CLASSES = list(MODEL_FOR_CAUSAL_LM_MAPPING.keys())
MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
@ -18,7 +22,7 @@ DataClassType = NewType("DataClassType", Any)
class ArgumentParser(HfArgumentParser):
def parse_yaml_and_args(
self, yaml_arg: str, other_args: Optional[list[str]] = None
self, yaml_arg: str, other_args: list[str] | None = None
) -> list[dataclass]:
"""
Parse a YAML file and overwrite the default/loaded values with the values provided to the command line.
@ -129,6 +133,10 @@ class ExperimentSetup(str, Enum):
@dataclass
class TrainingArguments(TrainingArguments):
tf32: bool = field(
default=True,
metadata={"help": "Whether to use tf32 precision."},
)
dataloader_pin_memory: bool = field(
default=True,
metadata={"help": "Whether to pin memory in data loaders or not."},
@ -143,11 +151,11 @@ class TrainingArguments(TrainingArguments):
},
)
dataloader_prefetch_factor: int = field(
default=2,
default=8,
metadata={"help": "Number of batches loaded in advance by each worker."},
)
dataloader_num_workers: int = field(
default=4,
default=8,
metadata={"help": "Number of subprocesses to use for data loading."},
)
optim: str = field(
@ -256,15 +264,15 @@ class ModelArguments:
@dataclass
class LoRAArguments:
lora_r: Optional[int] = field(
lora_r: int | None = field(
default=8,
metadata={"help": ("LoRA R value.")},
)
lora_dropout: Optional[float] = field(
lora_dropout: float | None = field(
default=0.02,
metadata={"help": ("LoRA dropout.")},
)
target_modules: Optional[list[str]] = field(
target_modules: list[str] | None = field(
default=None,
metadata={"help": ("LoRA target modules.")},
)
@ -284,11 +292,11 @@ class CtxTrainingArguments:
default=None,
metadata={"help": "Path to the pretrained checkpoint."},
)
max_base_len: Optional[int] = field(
max_base_len: int | None = field(
default=2**13,
metadata={"help": "Maximum base length for training."},
)
max_ctx_len: Optional[int] = field(
max_ctx_len: int | None = field(
default=2**13,
metadata={"help": "Maximum context length for training."},
)
@ -296,21 +304,21 @@ class CtxTrainingArguments:
default=False,
metadata={"help": "Whether to use sequence packing."},
)
per_device_train_max_batch_len: Optional[int] = field(
per_device_train_max_batch_len: int | None = field(
default=2**12,
metadata={
"help": "Maximum batch length for training. Only used with multipack sampler."
},
)
max_new_tokens: Optional[int] = field(
max_new_tokens: int | None = field(
default=2**10,
metadata={"help": "Maximum new tokens for generation-based evaluation."},
)
gen_per_device_eval_batch_size: Optional[int] = field(
gen_per_device_eval_batch_size: int | None = field(
default=1,
metadata={"help": "Per device evaluation batch size for generation."},
)
notes: Optional[str] = field(
notes: str | None = field(
default=None,
metadata={"help": "Wandb notes for the experiment."},
)
@ -342,19 +350,19 @@ class DataArguments:
default=False,
metadata={"help": "Whether to use streaming dataset for training."},
)
val_ds_names: Optional[list[str]] = field(
val_ds_names: list[str] | None = field(
default=None,
metadata={"help": "Validation dataset names."},
)
test_ds_names: Optional[list[str]] = field(
test_ds_names: list[str] | None = field(
default=None,
metadata={"help": "Test dataset names."},
)
max_val_samples_per_ds: Optional[int] = field(
max_val_samples_per_ds: int | None = field(
default=5000,
metadata={"help": "Maximum number of validation samples per dataset."},
)
max_test_samples_per_ds: Optional[int] = field(
max_test_samples_per_ds: int | None = field(
default=1000,
metadata={"help": "Maximum number of test samples per dataset."},
)
@ -378,7 +386,7 @@ class HypernetArguments:
default=0.0,
metadata={"help": "Dropout rate for HyperLoRA."},
)
extra_modules: Optional[list[str]] = field(
extra_modules: list[str] | None = field(
default=None,
metadata={"help": "Extra modules to train."},
)
@ -402,7 +410,7 @@ class CtxEncoderArguments:
default=None,
metadata={"help": "Context encoder model name or path."},
)
layer_idx: Optional[int] = field(
layer_idx: int | None = field(
default=None,
metadata={
"help": "Layer index for context encoder. "
@ -413,7 +421,6 @@ class CtxEncoderArguments:
@dataclass
class AggregatorArguments:
aggregator_type: Literal["pooler", "perceiver"] = field(
default="pooler",
metadata={"help": "Aggregator type for HyperLoRA."},

373
src/ctx_to_lora/data_definitions.py
View file

@ -1,373 +0,0 @@
from glob import glob
IGNORE_INDEX = -100
FW_QA_PATHS = [
f"data/raw_datasets/fw_qa/{i:05d}.parquet" for i in [0, 1, 6, 7, 8, 10, 22, 30, 35]
]
TRANSFORMED_DATA_DIR = "data/processed_datasets"
# approximate length of the datasets (train split)
# needed for streaming datasets
DS_LEN = {
"fw_qa_3_mini": 100_000,
"fw_qa_3_medium": 121_000_000,
"fw_qa_3": 270_000_000,
"fw_qa_xl": 27_000_000,
"ctx_qa": 300_000,
"pwc": 240_000,
"hotpot_qa": 100_000,
"squad": 90_000,
"drop": 77_000,
"narrativeqa": 40_000,
"quoref": 11_000,
"ropes": 11_000,
"synthetic_convqa": 40_000,
}
LONGBENCH_TASKS = [
"longbench/narrativeqa",
"longbench/qasper",
"longbench/multifieldqa_en",
"longbench/multifieldqa_zh",
"longbench/hotpotqa",
"longbench/2wikimqa",
"longbench/musique",
"longbench/dureader",
"longbench/gov_report",
"longbench/qmsum",
"longbench/multi_news",
"longbench/vcsum",
]
LONGBENCH_E_TASKS = [
"longbench/qasper_e",
"longbench/multifieldqa_en_e",
"longbench/hotpotqa_e",
"longbench/2wikimqa_e",
"longbench/gov_report_e",
"longbench/multi_news_e",
]
DS_KWARGS = {
"hotpot_qa": dict(
train=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="train[900:]"),
validation=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="train[:900]"),
test=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="validation"),
),
"hotpot_qa_tiny": dict(
train=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="train[900:2000]"),
validation=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="train[:900]"),
),
"pwc": dict(
train=dict(path="sggetao/PwC", split="train[900:]"),
validation=dict(path="sggetao/PwC", split="train[:900]"),
test=dict(path="sggetao/PwC", split="test"),
),
"pwc_tiny": dict(
train=dict(path="sggetao/PwC", split="train[900:2000]"),
validation=dict(path="sggetao/PwC", split="train[:900]"),
),
"squad": dict(
train=dict(path="rajpurkar/squad", split="train[900:]"),
validation=dict(path="rajpurkar/squad", split="train[:900]"),
test=dict(path="rajpurkar/squad", split="validation"),
),
# "fw_qa_tiny": dict(
# train=dict(
# path="parquet",
# data_files="data/raw_datasets/fw_qa/00000.parquet",
# split="train",
# ),
# validation=dict(
# path="parquet",
# data_files="data/raw_datasets/fw_qa/00000_val.parquet",
# split="train",
# ),
# ),
# "fw_qa": dict(
# train=dict(
# path="parquet",
# data_files=FW_QA_PATHS,
# split="train",
# ),
# validation=dict(
# path="parquet",
# data_files="data/raw_datasets/fw_qa/*_val.parquet",
# split="train",
# ),
# ),
# "fw_qa_large": dict(
# train=dict(
# path="parquet",
# data_files=glob("data/raw_datasets/fw_qa_large/*res.parquet"),
# split="train",
# ),
# validation=dict(
# path="parquet",
# data_files=glob("data/raw_datasets/fw_qa_large/*val.parquet"),
# split="train",
# ),
# ),
"fw_qa_xl": dict(
train=dict(
path="parquet",
data_files=glob("data/raw_datasets/fw_qa_xl/*[!val].parquet"),
split="train",
),
validation=dict(
path="parquet",
data_files=glob("data/raw_datasets/fw_qa_xl/*val.parquet"),
split="train",
),
# there is no explcit test for this ds
# used only for sanity check
test=dict(
path="parquet",
data_files=glob("data/raw_datasets/fw_qa_xl/*val.parquet"),
split="train",
),
),
# "fw_qa_2": dict(
# train=dict(
# path="parquet",
# data_files=glob("data/raw_datasets/fw_qa_2/*[!val].parquet"),
# split="train",
# ),
# validation=dict(
# path="parquet",
# data_files=glob("data/raw_datasets/fw_qa_2/*val.parquet"),
# split="train",
# ),
# ),
"fw_qa_3_mini": dict(
train=dict(
path="parquet",
data_files=glob("data/raw_datasets/fw_qa_3/000_00001.parquet"),
split="train[:100000]",
),
),
"fw_qa_3_medium": dict(
train=dict(
path="parquet",
data_files=glob("data/raw_datasets/fw_qa_3/00[0-5]*[!val].parquet"),
split="train",
),
),
"fw_qa_3": dict(
train=dict(
path="parquet",
data_files=glob("data/raw_datasets/fw_qa_3/*[!val].parquet"),
split="train",
),
validation=dict(
path="parquet",
data_files=glob("data/raw_datasets/fw_qa_3/*val.parquet"),
split="train",
),
),
"ctx_qa": dict(
train=dict(
path="parquet",
data_files=glob("data/raw_datasets/ctx_qa/*res.parquet"),
split="train",
),
validation=dict(
path="parquet",
data_files=glob("data/raw_datasets/ctx_qa/*val.parquet"),
split="train",
),
# there is no explcit test for this ds
# used only for sanity check
test=dict(
path="parquet",
data_files=glob("data/raw_datasets/ctx_qa/*val.parquet"),
split="train",
),
),
"drop": dict(
train=dict(
path="ucinlp/drop",
split="train",
),
),
"narrativeqa": dict(
train=dict(
path="nvidia/ChatQA-Training-Data",
name="narrativeqa",
split="train",
),
),
"quoref": dict(
train=dict(
path="nvidia/ChatQA-Training-Data",
name="quoref",
split="train",
),
),
"ropes": dict(
train=dict(
path="nvidia/ChatQA-Training-Data",
name="ropes",
split="train",
),
),
"synthetic_convqa": dict(
train=dict(
path="nvidia/ChatQA-Training-Data",
name="synthetic_convqa",
split="train",
),
),
"booksum": dict(
train=dict(
path="kmfoda/booksum",
split="train",
),
validation=dict(
path="kmfoda/booksum",
split="validation",
),
test=dict(
path="kmfoda/booksum",
split="test",
),
),
"gov_report": dict(
train=dict(
path="ccdv/govreport-summarization",
split="train",
),
validation=dict(
path="ccdv/govreport-summarization",
split="validation",
),
test=dict(
path="ccdv/govreport-summarization",
split="test",
),
),
# "wikitext-2": dict(
# train=dict(
# path="EleutherAI/wikitext_document_level",
# name="wikitext-2-raw-v1",
# split="train",
# ),
# ),
# "wikitext-103": dict(
# train=dict(
# path="EleutherAI/wikitext_document_level",
# name="wikitext-103-raw-v1",
# split="train",
# ),
# ),
"gsm8k": dict(
train=dict(
path="openai/gsm8k",
name="main",
split="train[100:]",
),
validation=dict(
path="openai/gsm8k",
name="main",
split="train[:100]",
),
),
"openmathintx-2": dict(
train=dict(
path="nvidia/OpenMathInstruct-2",
split="train_1M[:100000]",
),
),
"opencoder-edu": dict(
train=dict(
path="OpenCoder-LLM/opc-sft-stage2",
name="educational_instruct",
split="train[900:]",
),
validation=dict(
path="OpenCoder-LLM/opc-sft-stage2",
name="educational_instruct",
split="train[:900]",
),
),
"triviaqa_retrieved": dict(
test=dict(
path="json",
data_files="data/eval/triviaqa/triviaqa_retrieved.jsonl",
split="train",
),
),
# for negative training and evaluating
"negative_nq": dict(
test=dict(
path="json",
data_files="data/raw_datasets/negative_natural_questions/validation.jsonl",
split="train",
),
),
}
# LongBench kwargs
for ds_name in LONGBENCH_TASKS + LONGBENCH_E_TASKS:
DS_KWARGS[ds_name] = dict(
test=dict(
path="THUDM/LongBench",
name=ds_name.split("/")[-1],
split="test",
)
)
# for training closed qa datasets, e.g., hotpot_qa, squad, etc.
CLOSED_QA_INTX_TEMPLATES = [
"Answer the question based on the given passages. Only give me the answer and do not output any other words.\n\nQuestion: {input}",
"Answer without any explanation.\n\nQuestion: {input}",
"Based on the provided text, what is the answer to the following question? Provide only the answer.\n\nQuestion: {input}",
"Extract the answer to the question from the text. Be concise. Do not explain.\n\nQuestion: {input}",
"What is the answer to this question, based on the context? Respond with the answer only.\n\nQuestion: {input}",
"Provide a direct answer to the question using the given passages. Do not give any explanation.\n\nQuestion: {input}",
"Answer the question using only information from the provided text. No extra words.\n\nQuestion: {input}",
"From the passages, answer the question. Just the answer, please.\n\nQuestion: {input}",
"Give the answer to the question. Do not include any other text.\n\nQuestion: {input}",
"The answer to the question is in the text. Find it and state it clearly. No need for explanation.\n\nQuestion: {input}",
"Concisely answer the question based on the text provided. Don't include any other words. Just the answer.\n\nQuestion: {input}",
"Read the passages and answer the question with the minimal necessary words.\n\nQuestion: {input}",
"What is the direct response to the question, according to the text? Avoid explanation.\n\nQuestion: {input}",
"Please provide only the answer to the question, derived from the text.\n\nQuestion: {input}",
"Using the provided context, answer the question. Output the answer and nothing else.\n\nQuestion: {input}",
"Identify the answer in the text and present it without elaboration.\n\nQuestion: {input}",
"Answer the following question based on the text. Your answer should be brief and to the point. No explanation.\n\nQuestion: {input}",
"Based on the information given, what is the answer to the question? Only state the answer.\n\nQuestion: {input}",
"Find the answer to the question in the provided passages and write it down. No explanations.\n\nQuestion: {input}",
"The question is: {input}. Provide the answer based on the text, and nothing more.",
"Question: {input}\nAnswer directly based on the text provided. No extra words.",
"Question: {input}\nPlease provide the answer based on the text. No explanation is needed.",
]
EVAL_INTX_TEMPLATES = {
"negative_nq": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}",
"triviaqa_retrieved": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}",
"hotpot_qa": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}",
"squad": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}",
"longbench/qasper": 'Answer the question as concisely as you can, using a single phrase or sentence if possible.\nIf the question cannot be answered based on the information in the article, write "unanswerable".\nIf the question is a yes/no question, answer "yes", "no", or "unanswerable". Do not provide any explanation.\n\nQuestion: {input}',
"longbench/narrativeqa": "Answer the question as concisely as you can, using a single phrase if possible. Do not provide any explanation.\n\nQuestion: {input}",
"longbench/multifieldqa_en": "Answer the following question. Only output the answer and do not output any other words.\n\nQuestion: {input}",
# "longbench/multifieldqa_zh": "阅读以下文字并用中文简短回答：\n\n现在请基于上面的文章回答下面的问题，只告诉我答案，不要输出任何其他字词。\n\n问题：{input}",
"longbench/hotpotqa": "Answer the following question. Only output the answer and do not output any other words.\n\nQuestion: {input}",
"longbench/2wikimqa": "Answer the following question. Only output the answer and do not output any other words.\n\nQuestion: {input}",
"longbench/musique": "Answer the following question. Only output the answer and do not output any other words.\n\nQuestion: {input}",
#
# "longbench/dureader": "请基于给定的文章回答下述问题\n\n请基于上述文章回答下面的问题。\n\n问题：{input}",
"longbench/gov_report": "You are given a report by a government agency. Write a one-page summary of the report.",
"longbench/qmsum": "You are given a meeting transcript and a query containing a question or instruction. Answer the query based on the above meeting transcript.\n\nQuery: {input}",
"longbench/multi_news": "You are given several news passages. Write a one-page summary of all the news.",
# "longbench/vcsum": "下面有一段会议记录，请你阅读后，写一段总结，总结会议的内容。",
}
for ds_name in LONGBENCH_E_TASKS:
EVAL_INTX_TEMPLATES[ds_name] = EVAL_INTX_TEMPLATES[ds_name[:-2]]
for ds_name in DS_KWARGS:
if ds_name not in EVAL_INTX_TEMPLATES:
EVAL_INTX_TEMPLATES[ds_name] = "{input}"

980
src/ctx_to_lora/data_utils.py
View file

@ -1,980 +0,0 @@
import logging
import hashlib
import json
import random
from os import path
from typing import Any, Callable, Iterator, Optional
import datasets
import numpy as np
from datasets import load_dataset
from transformers import PreTrainedTokenizerBase
from ctx_to_lora.training_utils import TRAINING_TASK
from ctx_to_lora.data_definitions import (
DS_KWARGS,
EVAL_INTX_TEMPLATES,
TRANSFORMED_DATA_DIR,
IGNORE_INDEX,
CLOSED_QA_INTX_TEMPLATES,
)
logger = logging.getLogger()
# TODO: add continued pre-training pipeline
def closed_qa_prompting(prompt: str):
template = random.choice(CLOSED_QA_INTX_TEMPLATES)
return template.format(input=prompt)
# TODO: dont use closed_qa_prompting when training on self-generated responses
def get_preprocessing_fn(
ds_name: str, is_eval: bool
) -> Callable[[dict[str, Any]], dict[str, Any]]:
"""
Get preprocessing function for a specific dataset.
Args:
ds_name: Name of the dataset
Returns:
A preprocessing function that takes and returns a dictionary
"""
f = lambda x: x
if ds_name.startswith("pwc"):
def f(sample):
return {
"context": sample["input"],
"prompt": sample["prompt"],
"response": sample["answer"],
}
elif ds_name.startswith("hotpot_qa"):
def f(sample):
txt = ""
for p in sample["context"]["sentences"]:
txt += " " + "".join(p)
q = sample["question"]
prompt = closed_qa_prompting(q) if not is_eval else q
return {
"context": txt.strip(),
"prompt": prompt,
"response": sample["answer"],
}
elif "squad" in ds_name:
def f(sample):
q = sample["question"]
prompt = closed_qa_prompting(q) if not is_eval else q
return {
"context": sample["context"],
"prompt": prompt,
"response": sample["answers"]["text"][0],
}
elif ds_name == "drop":
def f(sample):
q = sample["question"]
prompt = closed_qa_prompting(q) if not is_eval else q
return {
"context": sample["passage"],
"prompt": prompt,
"response": ", ".join(sample["answers_spans"]["spans"]),
}
elif ds_name in ["narrativeqa", "quoref", "ropes"]:
def f(sample):
response = sample["answers"][0]
if isinstance(response, list):
response = response[0]
q = sample["messages"][-1]["content"]
prompt = closed_qa_prompting(q) if not is_eval else q
return {
"context": sample["document"],
"prompt": prompt,
"response": response,
}
elif ds_name == "synthetic_convqa":
def f(sample):
response = sample["answers"][0]
if isinstance(response, list):
response = response[0]
return {
"context": sample["document"],
"prompt": sample["messages"][-1]["content"],
"response": response,
}
elif ds_name == "booksum":
prompt_templates = [
"Summarization the provided text.",
"# Summary",
"### Summary",
"Summary of the text",
]
# TODO: use these templates for training
# analysis_templates = []
# summary_templates = []
def f(sample):
return {
"context": sample["chapter"],
"prompt": "Summarize the provided text.",
"response": sample["summary_text"],
}
elif ds_name == "gov_report":
def f(sample):
return {
"context": sample["report"],
"prompt": "Summarize the provided text.",
"response": sample["summary"],
}
elif "wikitext" in ds_name:
def f(sample):
return {
"context": sample["page"],
"prompt": "PLAECHOLDER",
"response": "PLAECHOLDER",
}
elif ds_name == "openmathintx-2":
def f(sample):
return {
"context": sample["problem"],
"prompt": sample["problem"],
"response": sample["generated_solution"],
}
elif "gsm8k" in ds_name:
def f(sample):
return {
"context": sample["question"],
"prompt": sample["question"],
"response": sample["answer"],
}
elif "opencoder-edu" in ds_name:
def f(sample):
return {
"context": sample["instruction"],
"prompt": sample["instruction"],
"response": "```python\n" + sample["code"].strip() + "\n```",
}
elif ds_name.startswith("longbench"):
def f(sample):
return {
"context": sample["context"],
"prompt": sample["input"],
"response": sample["answers"][0],
}
elif ds_name == "negative_nq":
def f(sample):
q = sample["prompt"]
prompt = closed_qa_prompting(q) if not is_eval else q
return {
"context": sample["context"],
"prompt": prompt,
"response": sample["answer"],
}
elif ds_name == "triviaqa_retrieved":
# only used for eval
def f(sample):
return {
"context": sample["context"],
"prompt": sample["prompt"],
"response": sample["answer"],
}
if is_eval and ds_name in EVAL_INTX_TEMPLATES:
prompt_template = EVAL_INTX_TEMPLATES[ds_name]
def eval_intx_decorator(f):
def g(sample):
sample = f(sample)
sample["prompt"] = prompt_template.format(input=sample["prompt"])
return sample
return g
return eval_intx_decorator(f) if is_eval else f
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: {kwargs}"
)
# return kwargs
else:
# return DS_KWARGS[ds_name][split]
kwargs = DS_KWARGS[ds_name][split]
# custom logic for slicing iterable datasets
take, skip = None, None
kw_split = kwargs["split"]
if ("[" in kw_split) and kw_split.endswith("]"):
kwargs["split"], slice = kw_split.split("[")
slice = slice.strip("]")
skip = slice.split(":")[0]
if skip:
kwargs["skip"] = int(skip)
take = slice.split(":")[1]
if take:
kwargs["take"] = int(take)
return kwargs
def validate_columns(tokenized_ds):
cols = ["input_ids", "attention_mask", "labels"]
if "ctx_ids" in tokenized_ds.column_names:
cols += ["ctx_ids", "ctx_attn_mask"]
if "chat_ids" in tokenized_ds.column_names:
cols += ["chat_ids", "chat_attn_mask", "chat_labels"]
ref_cols = set(cols)
assert set(tokenized_ds.column_names) == ref_cols, (
f"Columns mismatch: {set(tokenized_ds.column_names)} != {ref_cols}"
)
def filter_long_samples(samples):
return [len(ctx) < 10000 for ctx in samples["context"]]
def filter_long_chat(samples):
return [len(chat) < 2**13 for chat in samples["chat"]]
def add_repeat_prompt_fn(samples):
unique_contexts = set()
ctxs, prompts, responses = [], [], []
for ctx, prompt, response in zip(
samples["context"], samples["prompt"], samples["response"]
):
# Only process if the context is not already in the set
if ctx in unique_contexts:
continue
# remove too long contexts
if len(ctx) > 1000:
continue
unique_contexts.add(ctx)
ctxs.append(ctx)
responses.append(ctx)
prompts.append("Repeat the text above.")
logger.debug("Adding repeat prompt...")
logger.debug(f"# unique contexts: {len(unique_contexts)}")
return dict(
context=ctxs + samples["context"],
response=responses + samples["response"],
prompt=prompts + samples["prompt"],
)
def add_negative_prompt_fn(samples):
unique_contexts = set()
ctxs, prompts, responses = [], [], []
keywords = [
"repeat",
"rephrase",
"summarize",
"rewrite",
"title",
"keyword",
"continuation",
]
for ctx, prompt, response in zip(
samples["context"], samples["prompt"], samples["response"]
):
if ctx in unique_contexts:
continue
unique_contexts.add(ctx)
if len(ctx) > 3000:
continue
if any(keyword in prompt for keyword in keywords):
# Skip samples where the prompt contains any of the specified keywords
continue
ctxs.append(ctx)
prompts.append(prompt)
responses.append(response)
logger.debug("Adding negative prompt...")
logger.debug(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 dict(
context=neg_ctxs + samples["context"],
prompt=neg_prompts + samples["prompt"],
response=neg_responses + samples["response"],
)
def filter_none(samples):
out = [True] * len(samples["context"])
for k in samples:
for i, sample in enumerate(samples[k]):
if not bool(sample):
out[i] = False
return out
def _load_and_process_dataset(
ds_name: str,
split: str,
add_negative_prompt: bool,
add_repeat_prompt: bool,
streaming: bool,
ds_path: str,
num_proc: int,
):
try:
ds_kwargs = get_ds_kwargs(ds_name, split)
skip = ds_kwargs.pop("skip", None)
take = ds_kwargs.pop("take", None)
# only load as stream if the dataset is local (parquet)
load_as_stream = (
("fw_qa" in ds_name)
and ("tiny" not in ds_name)
and ("mini" not in ds_name)
and streaming
and ds_kwargs["path"] == "parquet"
)
ds = load_dataset(
**ds_kwargs,
trust_remote_code=True,
streaming=load_as_stream,
)
if split == "train" and streaming and (not load_as_stream):
# if the dataset is hosted on HF, we load it first then convert to iterable
ds = ds.to_iterable_dataset(num_shards=128)
if skip is not None:
ds = ds.skip(skip)
if take is not None:
ds = ds.take(take)
except ValueError as e:
logger.warning(
f"Failed to load dataset {ds_name} with split {split}. Error: {e}\nSkipping..."
)
return None
cols_to_remove = [
col for col in ds.column_names if col not in ["context", "prompt", "response"]
]
is_eval = split != "train"
ds = ds.map(
get_preprocessing_fn(ds_name, is_eval),
remove_columns=cols_to_remove,
num_proc=num_proc,
)
# ds = ds.remove_columns(cols_to_remove)
ds = ds.filter(filter_none, batched=True, num_proc=num_proc)
if split == "train":
# TODO: drop samples or split into multiple contexts
ds = ds.filter(filter_long_samples, batched=True, num_proc=num_proc)
if add_negative_prompt:
# TODO: make explicit dataset for this
ds = ds.map(
add_negative_prompt_fn,
batched=True,
batch_size=100_000,
num_proc=num_proc,
)
if add_repeat_prompt and "context_numbers" not in ds_name:
# TODO: convert to AE/LM
ds = ds.map(
add_repeat_prompt_fn,
batched=True,
batch_size=100_000,
num_proc=num_proc,
)
ds.save_to_disk(ds_path, num_proc=num_proc)
return ds
def get_tokenized_dataset(
ds_name: str,
split: str,
base_model_max_len: int,
tokenizer: PreTrainedTokenizerBase,
tokenizer_kwargs: dict[str, Any],
ctx_model_max_len: int,
ctx_tokenizer: PreTrainedTokenizerBase,
ctx_tokenizer_kwargs: dict[str, Any],
add_ctx_to_chat: bool,
add_repeat_prompt: bool,
add_negative_prompt: bool,
use_kl_loss: bool,
set_format: Optional[str] = None,
streaming: bool = False,
) -> dict[str, Any]:
# TODO: update hash function to include max_len
# TODO: max_len could be included in the kwargs
assert not use_kl_loss, "KL loss is deprecated"
logger.debug(f"Loading dataset {ds_name} with split {split}...")
need_ctx_ids = not add_ctx_to_chat and ctx_model_max_len
load_and_process_kwargs = dict(
ds_name=ds_name,
split=split,
add_negative_prompt=add_negative_prompt,
add_repeat_prompt=add_repeat_prompt,
streaming=streaming,
)
num_proc = None if streaming and split == "train" else 8
ds_hash = hashlib.sha256(json.dumps(load_and_process_kwargs).encode()).hexdigest()
ds_path = f"{TRANSFORMED_DATA_DIR}/{ds_hash}"
if path.exists(ds_path) and split == "train":
# load the cached ds
logger.info(f"Loaded processed dataset from {ds_path}")
else:
logger.info(f"Loading dataset {ds_name} with split {split}...")
ds = _load_and_process_dataset(
**load_and_process_kwargs,
ds_path=ds_path,
num_proc=num_proc,
)
if split == "train":
# force loading the cached version for newly created ds
ds = datasets.load_from_disk(ds_path)
tokenized_ds = construct_and_tokenize_ctx_qa(
base_model_max_len,
tokenizer,
tokenizer_kwargs,
ctx_model_max_len,
ctx_tokenizer,
ctx_tokenizer_kwargs,
add_ctx_to_chat,
use_kl_loss,
need_ctx_ids,
ds,
split,
set_format,
num_proc,
)
return tokenized_ds
def construct_and_tokenize_ctx_qa(
base_model_max_len,
tokenizer,
tokenizer_kwargs,
ctx_model_max_len,
ctx_tokenizer,
ctx_tokenizer_kwargs,
add_ctx_to_chat,
use_kl_loss,
need_ctx_ids,
ds,
split,
set_format=None,
num_proc=None,
):
is_eval = split != "train"
# TODO: update hash function to include max_len
# TODO: max_len could be included in the kwargs
kwargs = dict(
tokenizer=repr(tokenizer),
tokenizer_kwargs=json.dumps(tokenizer_kwargs),
ctx_tokenizer=repr(ctx_tokenizer),
ctx_tokenizer_kwargs=json.dumps(ctx_tokenizer_kwargs),
add_ctx_to_chat=add_ctx_to_chat,
use_kl_loss=use_kl_loss,
need_ctx_ids=need_ctx_ids,
ds=ds._fingerprint,
# TODO: add split (after moving to continued pre-training pipeline)
set_format=set_format,
)
kwargs_str = json.dumps(kwargs)
logger.debug(f"Tokenizing dataset with kwargs: {kwargs_str}")
ds_hash = hashlib.sha256(kwargs_str.encode()).hexdigest()
ds_path = f"{TRANSFORMED_DATA_DIR}/{ds_hash}"
if path.exists(ds_path) and split == "train":
# load the cached ds
logger.info(f"Loaded tokenized dataset from {ds_path}")
ds = datasets.load_from_disk(ds_path)
return ds
# for sft + chat_model, we need to convert the dataset to chat format
# add "messages" field
ds = ds.map(
convert_ctx_prompt_response_to_messages,
fn_kwargs={"add_ctx_to_chat": add_ctx_to_chat},
num_proc=num_proc,
)
# add "chat" field
ds = ds.map(
get_sft_prompt_formatting_fn(TRAINING_TASK.COMPLETION, tokenizer),
batched=True,
batch_size=100_000,
num_proc=num_proc,
)
# tokenize the chat + mask the assistant inputs
# TODO: drop samples or split into multiple contexts
if split == "train":
ds = ds.filter(
filter_long_chat,
batched=True,
batch_size=100_000,
num_proc=num_proc,
)
# add "input_ids", "attention_mask", "labels"
tokenized_ds = ds.map(
tokenize_chat_messages,
fn_kwargs={
"tokenizer": tokenizer,
"mask_assistant_inputs": True,
"tokenizer_kwargs": tokenizer_kwargs,
},
num_proc=num_proc,
)
if split == "train":
# drop?
# base model cant handle these samples naturally
# should skip
...
else:
tokenized_ds = tokenized_ds.map(
truncate_middle_if_too_long,
fn_kwargs={
"max_length": base_model_max_len,
"columns": ["input_ids", "attention_mask"],
"max_new_tokens": 256,
},
batched=True,
batch_size=100_000,
num_proc=num_proc,
)
tokenized_ds = tokenized_ds.map(
add_length_info,
fn_kwargs={"columns": ["input_ids"]},
batched=True,
batch_size=100_000,
num_proc=num_proc,
)
# for use_kl_loss, we need "chat_ids" and "chat_attn_mask"
if use_kl_loss:
raise NotImplementedError("KL loss deprecated")
tokenized_ds = tokenized_ds.map(
convert_ctx_prompt_response_to_messages,
fn_kwargs={"add_ctx_to_chat": True},
remove_columns=["messages"],
)
tokenized_ds = tokenized_ds.map(
get_sft_prompt_formatting_fn(TRAINING_TASK.COMPLETION, tokenizer),
remove_columns=["chat"],
)
tokenized_ds = tokenized_ds.map(
tokenize_chat_messages,
fn_kwargs={
"tokenizer": tokenizer,
"mask_assistant_inputs": True,
"tokenizer_kwargs": tokenizer_kwargs,
"for_kl_loss": True,
},
)
if need_ctx_ids:
# tokenize the ctx_text to get ctx_ids and ctx_attn_mask
tokenized_ds = tokenized_ds.map(
tokenize_ctx_text,
fn_kwargs={"tokenizer": ctx_tokenizer},
batched=True,
batch_size=100_000,
num_proc=num_proc,
)
if split == "train":
# TODO: do something if ctx length is longer than the ctx model length
# e.g., drop or split for multi-lora training
...
else:
# truncate in the middle
tokenized_ds = tokenized_ds.map(
truncate_middle_if_too_long,
fn_kwargs={
"max_length": ctx_model_max_len,
"columns": ["ctx_ids", "ctx_attn_mask"],
# no need to add new_tokens here (used only for the ctx encoder)
"max_new_tokens": 0,
},
batched=True,
batch_size=100_000,
num_proc=num_proc,
)
tokenized_ds = tokenized_ds.map(
add_length_info,
fn_kwargs={"columns": ["ctx_ids"]},
batched=True,
batch_size=100_000,
num_proc=num_proc,
)
tokenized_ds = tokenized_ds.remove_columns(
["messages", "chat", "context", "prompt", "response"]
)
if set_format:
tokenized_ds.set_format(type=set_format)
# if isinstance(tokenized_ds, IterableDataset):
# # the columns are unknown when using streaming dataset
# tokenized_ds = tokenized_ds._resolve_features()
# validate_columns(tokenized_ds)
if split == "train":
tokenized_ds.save_to_disk(ds_path, num_proc=num_proc)
del tokenized_ds
return datasets.load_from_disk(ds_path)
return tokenized_ds
def get_sft_prompt_formatting_fn(
sft_mode: TRAINING_TASK,
tokenizer: PreTrainedTokenizerBase,
) -> Callable[[dict[str, Any]], dict[str, Any]]:
"""
Get a function that formats examples for supervised fine-tuning.
Args:
sft_mode: The training task type
tokenizer: The tokenizer to use for chat template application
Returns:
A function that takes a training example and returns formatted data
Raises:
NotImplementedError: If sft_mode is not COMPLETION or tokenizer has no chat template
"""
if sft_mode != TRAINING_TASK.COMPLETION:
raise NotImplementedError(
f"Training task {sft_mode} not supported. "
"Only completion is supported for now."
)
if tokenizer.chat_template is None:
raise NotImplementedError(
"Only chat models + SFT are supported. "
"Training with pre-training data is not supported yet."
)
# TODO: add support for causal_lm training (for pre-training data)
# TODO: add support for recon training (for pre-training data)
# TODO: add support for non-chat models
# def f(example):
# output_texts = (
# dict(text=[]) if sft_mode == "causal_lm" else dict(prompt=[], response=[])
# )
# df = pd.DataFrame(dict(example))
# for i, inp_txt in df.iterrows():
# if sft_mode == "causal_lm":
# text = metadata["text_template"].format(**inp_txt)
# output_texts["text"].append(text)
# elif sft_mode == "completion":
# prompt = metadata["user_prompt_template"].format(**inp_txt)
# output_texts["prompt"].append(prompt)
# output_texts["response"].append(str(inp_txt[metadata["response_field"]]))
# return output_texts
def f_intx(example):
chat_text = tokenizer.apply_chat_template(
example["messages"], tokenize=False, add_generation_prompt=False
)
return dict(chat=chat_text)
# return f if not apply_chat_template_fn is not None else f_intx
return f_intx
def convert_ctx_prompt_response_to_messages(
example: dict[str, Any],
add_ctx_to_chat: bool = True,
) -> dict[str, Any]:
"""
Convert context/prompt/response format to chat messages format.
Args:
example: Dictionary containing 'prompt' and 'response' keys
add_ctx_to_chat: Whether to prepend context to the user message
Returns:
Dictionary with added 'messages' key containing chat format
Raises:
ValueError: If 'prompt' or 'response' keys are missing
"""
if "prompt" not in example or "response" not in example:
raise ValueError(f"'prompt' and 'response' are required. Got: {example}")
system_msg = ""
if "system_message" in example:
system_msg = example["system_message"].strip()
user_msg = example["prompt"].strip()
if add_ctx_to_chat:
user_msg = example["context"].strip() + "\n\n" + user_msg.strip()
messages = [
{"role": "system", "content": system_msg.strip()},
{"role": "user", "content": user_msg.strip()},
{"role": "assistant", "content": example["response"].strip()},
]
return dict(messages=messages)
# adapted from https://github.com/huggingface/trl/issues/632#issuecomment-1972630547
def get_assistant_start_end_indices(
messages: list[dict[str, str]],
tokenizer: PreTrainedTokenizerBase,
conversation_ids: list[int],
) -> list[tuple[int, int]]:
"""
Get the start and end indices of assistant messages in conversation text using tokenized messages.
Args:
messages: List of message dictionaries with 'role' and 'content' keys
tokenizer: Tokenizer used for tokenization
conversation_ids: Tokenized conversation
Returns:
List of (start, end) index tuples for assistant messages in the tokenized conversation
"""
assistant_ranges = []
current_idx = 0
for idx, message in enumerate(messages):
# if message["role"] == "assistant":
# Tokenize the assistant message content
tokenized_message = tokenizer.encode(
message["content"], add_special_tokens=False
)
# Find the start index of the assistant message in the tokenized conversation
for i in range(current_idx, len(conversation_ids)):
if conversation_ids[i : i + len(tokenized_message)] == tokenized_message:
start_idx = i
end_idx = i + len(tokenized_message)
current_idx = end_idx
if message["role"] == "assistant":
if idx == (len(messages) - 1):
# include everything after the last assistant message
assistant_ranges.append((start_idx, len(conversation_ids)))
else:
assistant_ranges.append((start_idx, end_idx))
break
return assistant_ranges
def get_masked_labels(
conversation_ids: dict[str, list[Any]],
assistant_ranges: list[tuple[int, int]],
) -> Iterator[int]:
"""
Generate masked labels for conversation, masking non-assistant tokens.
NOTE: This will also includes extra tokens between assistant and user messages
in multi-turn conversations.
E.g., {assistant_msg} <|eot_id|><|start_header_id|>user<|end_header_id|> {user_msg}
for Llama 3 models.
Args:
conversation_ids: Dictionary with tokenized conversation info
assistant_ranges: List of (start, end) indices for assistant messages
Yields:
Token ID or IGNORE_INDEX for each position
"""
for id_, (id_s, id_e) in list(
zip(conversation_ids["input_ids"], conversation_ids["offset_mapping"])
):
if any(id_s >= s and id_e <= e for s, e in assistant_ranges):
yield id_
else:
yield IGNORE_INDEX
def tokenize_chat_messages(
example: dict[str, Any],
tokenizer: PreTrainedTokenizerBase,
mask_assistant_inputs: bool = True,
for_kl_loss: bool = False,
tokenizer_kwargs: Optional[dict[str, Any]] = None,
) -> dict[str, list[int]]:
"""
Tokenize chat messages and optionally mask non-assistant tokens.
Args:
example: Dictionary containing 'chat' and 'messages' keys
tokenizer: Tokenizer to use
mask_assistant_inputs: Whether to mask non-assistant tokens
for_kl_loss: change the column names to "chat_ids" and "chat_attn_mask"
tokenizer_kwargs: Additional arguments to pass to tokenizer
Returns:
Dictionary containing tokenized inputs and labels
"""
# should be used only with chat models
text = example["chat"]
messages = example["messages"]
n_response = len([m for m in messages if m["role"] == "assistant"])
if n_response != 1:
raise ValueError(f"Expected 1 assistant response. Got {n_response}.")
conversation_ids = tokenizer(
text,
# return_offsets_mapping=mask_assistant_inputs,
add_special_tokens=False,
truncation=False,
**(tokenizer_kwargs or {}),
)
if tokenizer_kwargs and (
len(conversation_ids["input_ids"]) >= tokenizer_kwargs["max_length"]
):
raise ValueError(
f"Conversation length {len(conversation_ids['input_ids'])} exceeds max length {tokenizer_kwargs['max_length']}"
)
if mask_assistant_inputs:
assistant_ranges = get_assistant_start_end_indices(
messages, tokenizer, conversation_ids["input_ids"]
)
# labels = get_masked_labels(conversation_ids, assistant_ranges)
labels = [
id_ if any(s <= i < e for s, e in assistant_ranges) else IGNORE_INDEX
for i, id_ in enumerate(conversation_ids["input_ids"])
]
conversation_ids["labels"] = labels
# del conversation_ids["offset_mapping"]
else:
conversation_ids["labels"] = conversation_ids["input_ids"]
if for_kl_loss:
conversation_ids["chat_ids"] = conversation_ids.pop("input_ids")
conversation_ids["chat_attn_mask"] = conversation_ids.pop("attention_mask")
conversation_ids["chat_labels"] = conversation_ids.pop("labels")
return conversation_ids
def add_length_info(samples: dict[str, any], columns: list[str]) -> dict[str, any]:
for col in columns:
samples[f"{col}_len"] = [len(t) for t in samples[col]]
return samples
def truncate_middle_if_too_long(
samples: dict[str, any],
max_length: int,
columns: list[str],
max_new_tokens: int = 256,
) -> dict[str, any]:
"""
Truncate the middle of a list of tokens to fit within a maximum length.
Args:
tokens: List of token IDs
max_length: Maximum length for the truncated tokens
Returns:
List of truncated token IDs
"""
max_new_tokens_half = max_new_tokens // 2
# leave max_new_tokens for generation
half = max_length // 2 - max_new_tokens_half
for col in columns:
# if len(samples[col]) <= max_length:
# return samples[col]
samples[col] = [
t[:half] + t[-half:] if len(t) > max_length else t for t in samples[col]
]
return samples
def tokenize_ctx_text(
samples: dict[str, Any],
tokenizer: PreTrainedTokenizerBase,
) -> dict[str, Any]:
if tokenizer.chat_template:
txt = tokenizer.apply_chat_template(
[[{"role": "user", "content": text.strip()}] for text in samples["context"]],
tokenize=False,
add_generation_prompt=False,
)
# special tokens are already added by the chat template
tokenized_text = tokenizer(txt, truncation=False, add_special_tokens=False)
else:
txt = samples["context"]
tokenized_text = tokenizer(txt, truncation=False)
ctx_ids = tokenized_text["input_ids"]
ctx_attn_mask = tokenized_text["attention_mask"]
return dict(ctx_ids=ctx_ids, ctx_attn_mask=ctx_attn_mask)
if __name__ == "__main__":
from transformers import AutoTokenizer
# model_name = "meta-llama/Llama-3.2-1B-Instruct"
# model_name = "meta-llama/Llama-2-7b-chat-hf"
model_name = "google/gemma-2-2b-it"
messages = [
{"role": "user", "content": "Hello!"},
{"role": "assistant", "content": "Hello!"},
# {"role": "user", "content": "Not too bad"},
# {"role": "assistant", "content": "Cooooooool"},
]
tokenizer = AutoTokenizer.from_pretrained(model_name)
chat = tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=False, add_special_tokens=False
)
print(chat)
model_inputs = tokenize_chat_messages(
{"chat": chat, "messages": messages},
tokenizer,
# for_kl_loss=True,
)
print(tokenizer(chat, add_special_tokens=False))
print(model_inputs)

22
src/ctx_to_lora/hooks.py
View file

@ -1,6 +1,6 @@
import logging
from collections.abc import Callable, Iterable
from operator import attrgetter
from typing import Callable, Iterable, Optional
import torch
import torch.nn.functional as F
@ -8,6 +8,7 @@ from einops import einsum
from jaxtyping import Float, Integer
from torch import Tensor
from torch.utils.hooks import RemovableHandle
from ctx_to_lora.utils import get_layers
logger = logging.getLogger()
@ -28,8 +29,8 @@ def remove_hook_handles(handles: list[RemovableHandle]) -> None:
def apply_hook_to_model(
model: torch.nn.Module,
module_name: str,
pre_hook: Optional[Callable] = None,
post_hook: Optional[Callable] = None,
pre_hook: Callable | None = None,
post_hook: Callable | None = None,
) -> torch.nn.Module:
if not ((pre_hook is not None) or (post_hook is not None)):
raise ValueError("No hooks provided. Nothing to apply.")
@ -46,9 +47,9 @@ def apply_hook_to_model(
def apply_hook_to_layer(
layer: torch.nn.Module,
mname: str,
pre_hook: Optional[Callable] = None,
post_hook: Optional[Callable] = None,
) -> tuple[Optional[RemovableHandle], Optional[RemovableHandle]]:
pre_hook: Callable | None = None,
post_hook: Callable | None = None,
) -> tuple[RemovableHandle | None, RemovableHandle | None]:
"""
Applies pre and/or post hooks to a specific layer in the model.
@ -85,8 +86,8 @@ def apply_hook_to_layers(
model: torch.nn.Module,
module_names: list[str],
layer_indices: Iterable[int],
pre_hook: Optional[Callable] = None,
post_hook: Optional[Callable] = None,
pre_hook: Callable | None = None,
post_hook: Callable | None = None,
) -> list[RemovableHandle]:
"""
Applies custom hooks to specified layers and modules in the model.
@ -133,7 +134,7 @@ def add_generated_lora_hook(
B: Float[Tensor, "bs d_out r"],
scaling: float,
input_dropout: float,
position_ids: Optional[Integer[Tensor, "bs seq_len"]] = None,
position_ids: Integer[Tensor, "bs seq_len"] | None = None,
training: bool = False,
) -> list[RemovableHandle]:
"""
@ -157,7 +158,6 @@ def add_generated_lora_hook(
args: tuple | Float[Tensor, "bs seq_len d_in"],
output: Float[Tensor, "bs seq_len d_out"],
) -> Float[Tensor, "bs seq_len d_out"]:
if isinstance(output, tuple):
model_out = output[0]
else:
@ -238,7 +238,7 @@ def add_generated_layernorm_hook(
layer_index: int,
W: Float[Tensor, "bs hidden_size"],
training: bool,
position_ids: Optional[Integer[Tensor, "bs seq_len"]] = None,
position_ids: Integer[Tensor, "bs seq_len"] | None = None,
) -> list[RemovableHandle]:
"""
Adds layer normalization hooks to specified modules and layers.

89
src/ctx_to_lora/metrics.py Normal file
View file

@ -0,0 +1,89 @@
from collections import defaultdict
from collections.abc import Callable
import numpy as np
import torch
from transformers import EvalPrediction
def compute_per_token_acc(shift_logits, shift_labels, valid_masks):
indices = torch.where(valid_masks)
acc = (shift_logits.argmax(-1) == shift_labels)[indices].float()
return {
"per_token_accs": acc.flatten().tolist(),
"n_per_token_accs": valid_masks.sum().item(),
}
def compute_prefix_matching(shift_logits, shift_labels, valid_masks):
lengths = valid_masks.sum(dim=1)
is_wrong = (shift_logits.argmax(-1) != shift_labels) * valid_masks
is_correct = (shift_logits.argmax(-1) == shift_labels) * valid_masks
# NOTE: not reliable for multi-turn conversations
# ie, all tokens in the following user's turn will be correct
# still monotonically correlate with perf though
wrong_pos = torch.argmax(is_wrong, dim=1) - torch.argmax(valid_masks, dim=1)
perf = wrong_pos / lengths
# if all tokens are correct, set to 1
perf = torch.where(is_correct.sum(dim=1) == lengths, 1, perf)
return {
"prefix_matchings": perf.tolist(),
"n_prefix_matchings": valid_masks.shape[0],
}
@torch.no_grad()
def compute_perplexity(shift_logits, shift_labels, valid_masks):
loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
loss = loss_fct(shift_logits.transpose(1, 2), shift_labels)
loss = (loss * valid_masks).sum(dim=1) / valid_masks.sum(dim=1)
preplexities = torch.exp(loss)
return {
"perplexities": preplexities.tolist(),
"n_perplexities": valid_masks.shape[0],
}
class Evaluator:
def __init__(self, metric_fns: list[Callable]):
self.metric_fns = metric_fns
self.reset()
def reset(self):
self.accum_metrics = defaultdict(list)
self.count = defaultdict(list)
def update(self, shift_logits, shift_labels, valid_masks):
for metric_fn in self.metric_fns:
metric = metric_fn(shift_logits, shift_labels, valid_masks)
for k, v in metric.items():
if k.startswith("n_"):
self.count[k[2:]].append(v)
else:
self.accum_metrics[k] += v
def compute(self):
# Get result across entire eval set
result = {
k: np.sum(v) / np.sum(self.count[k]) for k, v in self.accum_metrics.items()
}
# Reset batch statistics
self.reset()
return result
@torch.no_grad()
def compute_metrics(
eval_pred: EvalPrediction,
compute_result: bool,
evaluator: Evaluator,
) -> dict | None:
logits, labels = eval_pred.predictions, eval_pred.label_ids
shift_logits = logits[..., :-1, :]
shift_labels = labels[..., 1:]
valid_masks = torch.where(shift_labels != -100, 1, 0)
evaluator.update(shift_logits, shift_labels, valid_masks)
if compute_result:
return evaluator.compute()

69
src/ctx_to_lora/model_loading.py
View file

@ -2,7 +2,7 @@ import logging
import os
import torch
from peft import LoraConfig, PeftConfig, PeftModel, VeraConfig
from peft import PeftModel
from peft import get_peft_config as _get_peft_config
from peft.utils import PeftType
from transformers import (
@ -46,16 +46,9 @@ def get_model_and_tokenizer(
def get_tokenizer(
model_name_or_path, tokenizer_kwargs=None, peft_config=None, train=False
):
# tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, padding_side="left")
# NOTE: lora models don't have tokenizer config in the folder
# left pad for generation, right pad for training (why?)
padding_side = "left" if not train else "right"
truncation_side = "left"
# tokenizer = AutoTokenizer.from_pretrained(
# "models/Mistral-7B-v0.1/", padding_side=padding_side
# )
if peft_config:
model_name_or_path = peft_config.base_model_name_or_path
@ -72,10 +65,6 @@ def get_tokenizer(
if tokenizer.pad_token_id is None:
tokenizer.pad_token_id = tokenizer.eos_token_id
# tokenizer.pad_token = tokenizer.eos_token
# tokenizer.pad_token = tokenizer.eos_token
# tokenizer.add_bos_token = True
template_path = f"chat_templates/{model_name_or_path}.jinja"
if os.path.exists(template_path):
@ -83,26 +72,6 @@ def get_tokenizer(
chat_template = open(template_path).read()
chat_template = chat_template.replace(" ", "").replace("\n", "")
tokenizer.chat_template = chat_template
# assert os.path.exists(f"{model_name_or_path}/chat_template.jinja"), (
# f"Chat template not found in {model_name_or_path}\n\n"
# "We assume a specfic form of chat template for consistency between models. Please use the templates provided",
# )
# if os.path.exists(f"{model_name_or_path}/chat_template.jinja"):
# chat_template = open(f"{model_name_or_path}/chat_template.jinja").read()
# chat_template = chat_template.replace(" ", "").replace("\n", "")
# tokenizer.chat_template = chat_template
# tokenizer.add_eos_token = False
# if train:
# # NOTE: this correctly add an eos token that has attention = 1
# # while padding eos tokens have attention = 0
# tokenizer.add_eos_token = True
# # shouldn't be needed as we manually compute the labels now
# # NOTE: this seems oddly important
# # and is not necessarily in the alignment handbook scripts
# # tokenizer.pad_token = tokenizer.unk_token # fix model not generating eos
return tokenizer
@ -123,8 +92,6 @@ def get_model(
trust_remote_code=True,
attn_implementation="eager",
use_cache=False,
# load_in_4bit=True,
# load_in_8bit=True,
)
is_vision_model = "Llama" in model_name_or_path and "Vision" in model_name_or_path
if model_kwargs is not None:
@ -148,8 +115,6 @@ def get_model(
model_init_kwargs["torch_dtype"] = torch.float32
model_init_kwargs.pop("use_cache")
# if train and not is_vision_model:
# model_init_kwargs["use_cache"] = False
logger.debug(f"Model init kwargs: {model_init_kwargs}")
if not is_vision_model:
if is_bidir_model:
@ -159,21 +124,10 @@ def get_model(
else:
model = MllamaForConditionalGeneration.from_pretrained(**model_init_kwargs)
model = model.language_model
# model = AutoModelForCausalLM.from_pretrained(**model_init_kwargs)
if peft_config is not None:
model = PeftModel(model, peft_config)
model.train(train)
for name, param in model.named_parameters():
# if "modules_to_save" not in name:
# param.requires_grad = requires_grad
# else:
# logging.debug(f"modules_to_save: {name} to be trained")
# # always train "modules_to_save"
# if not "layernorm" in name:
# raise NotImplementedError(
# f"modules_to_save should only be layernorm, got {name}"
# )
# param.requires_grad = True
param.requires_grad = requires_grad
return model
@ -195,24 +149,3 @@ def get_lora_config(model_dir, **kwargs):
peft_conf_kwargs.update(kwargs)
peft_config = _get_peft_config(peft_conf_kwargs)
return peft_config
# def get_emb_model_and_fns(emb_model_name, device):
# emb_model = AutoModel.from_pretrained(
# emb_model_name,
# device_map=device,
# torch_dtype=torch.float32 if "gte" in emb_model_name else torch.bfloat16,
# trust_remote_code=True,
# ).eval()
# emb_tokenizer = AutoTokenizer.from_pretrained(emb_model_name)
# if emb_tokenizer.pad_token_id is None:
# emb_tokenizer.pad_token_id = emb_tokenizer.eos_token_id
# # assert bool(args.query_intx), "query_intx must be provided for the emb_model_name"
# # partial(get_detailed_instruct, task_description=args.query_intx)
# task_desc_format_fn = add_full_stop
# if "SFR" in emb_model_name:
# task_desc_format_fn = apply_sfr_template
# pooling_fn = get_pooling_fn("last_token")
# elif "gte" in emb_model_name:
# pooling_fn = get_pooling_fn("cls")
# return emb_model, emb_tokenizer, task_desc_format_fn, pooling_fn

0
src/ctx_to_lora/eval_utils.py → src/ctx_to_lora/modeling/__init__.py
View file

234
src/ctx_to_lora/modeling/aggregator.py Normal file
View file

@ -0,0 +1,234 @@
import logging
from dataclasses import dataclass
from enum import Enum
from einops import rearrange, unpack
from jaxtyping import Float, Integer
from torch import Tensor, nn
from transformers import (
PretrainedConfig,
PreTrainedModel,
)
from ctx_to_lora.configs import (
AggregatorArguments,
)
from ctx_to_lora.modeling.idefics2 import Idefics2Perceiver, Idefics2PerceiverConfig
from ctx_to_lora.pooling import POOL_FN
from ctx_to_lora.utils import (
get_num_layers,
)
logger = logging.getLogger()
class AGGREGATOR_TYPE(str, Enum):
POOLER = "pooler"
PERCEIVER = "perceiver"
@dataclass
class AggregatorConfig:
aggregator_type: AGGREGATOR_TYPE
num_layers: int
num_modules: int
num_extra_modules: int
output_size: int
feature_size: int
# pooler
pooling_type: POOL_FN
# perceiver
num_self_attends_per_block: int # = 16
decoder_depth: int # 1 = only cross-attention
num_latent_factor: int = 8
lora_r: int = 8
per_rank_gen: bool = False
def get_aggregator_config(
model: PreTrainedModel,
ctx_encoder_model_config: PretrainedConfig,
output_size: int,
num_modules: int,
num_extra_modules: int,
lora_r: int,
per_rank_gen: bool,
aggregator_args: AggregatorArguments,
):
return AggregatorConfig(
feature_size=ctx_encoder_model_config.hidden_size,
output_size=output_size,
num_layers=get_num_layers(model),
num_modules=num_modules,
num_extra_modules=num_extra_modules,
lora_r=lora_r,
per_rank_gen=per_rank_gen,
**vars(aggregator_args),
)
class Perceiver(nn.Module):
"""perceiver w/ bottleneck size = n_modules * n_layers"""
def __init__(
self,
feature_size,
output_size,
num_layers,
num_modules,
num_extra_modules,
per_rank_gen,
lora_r,
num_latent_factor,
*args,
**kwargs,
):
super().__init__()
self.num_layers = num_layers
self.num_modules = num_modules
self.num_extra_modules = num_extra_modules
self.per_rank_gen = per_rank_gen
self.r = lora_r if self.per_rank_gen else 1
n_output_queries = num_layers * (num_modules * self.r + num_extra_modules)
self.config = Idefics2PerceiverConfig(
input_size=feature_size,
# the first layer is xattn
resampler_depth=kwargs["num_self_attends_per_block"] + 1,
resampler_n_latents=n_output_queries * num_latent_factor,
intermediate_size_factor=4,
hidden_size=output_size,
attn_implementation="flash_attention_2",
)
self.decoder_config = Idefics2PerceiverConfig(
input_size=output_size,
resampler_depth=kwargs["decoder_depth"],
resampler_n_latents=n_output_queries,
intermediate_size_factor=4,
hidden_size=output_size,
attn_implementation="flash_attention_2",
)
self.perceiver = Idefics2Perceiver(self.config, self.decoder_config)
def forward(
self,
ctx_features: Float[Tensor, "bs seq_len feature_dim"],
ctx_attn_mask: Integer[Tensor, "bs seq_len"] | None = None,
ctx_position_ids: Integer[Tensor, "bs seq_len"] | None = None,
):
x = self.perceiver(ctx_features, ctx_attn_mask, ctx_position_ids)
lora_x, extra_x = unpack(
x,
[
[self.num_layers * self.num_modules * self.r],
[self.num_layers * self.num_extra_modules],
],
"bs * feature_dim",
)
lora_x = rearrange(
lora_x,
"bs (n_layers n_modules r) d -> bs n_layers n_modules r d",
n_modules=self.num_modules,
n_layers=self.num_layers,
r=self.r,
)
if not self.per_rank_gen:
lora_x = lora_x.squeeze(3)
extra_x = rearrange(
extra_x,
"bs (n_layers n_extra_modules) d -> bs n_layers n_extra_modules d",
n_extra_modules=self.num_extra_modules,
n_layers=self.num_layers,
)
# x = rearrange(
# x,
# "bs (n_layers n_modules) d -> bs n_layers n_modules d",
# n_modules=self.num_modules,
# n_layers=self.num_layers,
# )
# lora_emb, extra_emb = unpack(
# emb,
# [[self.num_modules], [self.num_extra_modules]],
# "bs n_layers * feature_dim",
# )
return lora_x, extra_x
# class Pooler(nn.Module):
# def __init__(
# self,
# feature_size: int,
# output_size: int,
# pooling_type: POOL_FN,
# num_layers: int,
# num_modules: int,
# *args,
# **kwargs,
# ):
# super().__init__()
# self.num_layers = num_layers
# self.num_modules = num_modules
# # NOTE: features will be projected to size = output_size // 2
# # then cat with layer and module embeddings (each with size output_size // 4)
# # which are collectively form features with size = output_size
# self.pool_fn = get_pooling_fn(pooling_type)
# self.feature_proj = nn.Linear(feature_size, output_size // 2)
# self.ln = nn.LayerNorm(output_size // 2)
# self.layer_embs = nn.Sequential(
# nn.Embedding(num_layers, output_size // 4),
# nn.LayerNorm(output_size // 4),
# )
# self.module_embs = nn.Sequential(
# nn.Embedding(num_modules, output_size // 4),
# nn.LayerNorm(output_size // 4),
# )
# self.mixer = Mixer(output_size, output_size * 4, output_size)
# self.mlp = MLPResidualBlock(output_size, output_size * 4, output_size)
# self.register_buffer("layer_indices", torch.arange(num_layers))
# self.register_buffer("module_indices", torch.arange(num_modules))
# def forward(
# self,
# features: Float[Tensor, "bs seq_len feature_dim"],
# attn_mask: Integer[Tensor, "bs seq_len"] | None = None,
# ):
# bs = features.shape[0]
# # [bs, feature_dim]
# x = self.ln(self.feature_proj(self.pool_fn(features, attn_mask).float()))
# x = repeat(
# x,
# "bs d -> bs n_layers n_modules d",
# n_modules=self.num_modules,
# n_layers=self.num_layers,
# )
# layer_embs = self.layer_embs(self.layer_indices) # [num_layers, d]
# layer_embs = repeat(
# layer_embs,
# "n_layers d -> bs n_layers n_modules d",
# bs=bs,
# n_modules=self.num_modules,
# )
# module_embs = self.module_embs(self.module_indices) # [num_modules, d]
# module_embs = repeat(
# module_embs,
# "n_modules d -> bs n_layers n_modules d",
# bs=bs,
# n_layers=self.num_layers,
# )
# emb = torch.cat([x, layer_embs, module_embs], dim=3)
# return self.mlp(self.mixer(emb))
AG = {
# AGGREGATOR_TYPE.POOLER: Pooler,
AGGREGATOR_TYPE.PERCEIVER: Perceiver,
}

65
src/ctx_to_lora/modeling/ctx_encoder.py Normal file
View file

@ -0,0 +1,65 @@
import logging
from contextlib import contextmanager
import torch
from torch import nn
from transformers import PreTrainedModel
logger = logging.getLogger()
@contextmanager
def early_exit(base_model: PreTrainedModel, exit_layer: int):
try:
layers = base_model.layers
base_model.layers = layers[:exit_layer]
yield base_model
finally:
base_model.layers = layers
@contextmanager
def maybe_add_batch_dim(kwargs):
try:
batched_input = False
batched_attn_mask = False
if (
"input_ids" in kwargs
and kwargs["input_ids"] is not None
and len(kwargs["input_ids"].shape) == 1
):
kwargs["input_ids"] = kwargs["input_ids"].unsqueeze(0)
batched_input = True
if (
"attention_mask" in kwargs
and kwargs["attention_mask"] is not None
and isinstance(kwargs["attention_mask"], torch.Tensor)
and len(kwargs["attention_mask"].shape) == 1
):
kwargs["attention_mask"] = kwargs["attention_mask"].unsqueeze(0)
batched_attn_mask = True
yield batched_input, batched_attn_mask
finally:
if batched_input:
kwargs["input_ids"] = kwargs["input_ids"].squeeze(0)
if batched_attn_mask:
kwargs["attention_mask"] = kwargs["attention_mask"].squeeze(0)
class EarlyExit(nn.Module):
def __init__(self, base_model: PreTrainedModel, exit_layer: int):
super().__init__()
self.base_model = base_model
if "gte" in base_model.config.name_or_path:
self.base_model.encoder.layer = base_model.encoder.layer[:exit_layer]
else:
self.base_model.layers = base_model.layers[:exit_layer]
@property
def config(self):
return self.base_model.config
@torch.no_grad()
def forward(self, **kwargs):
model_outputs = self.base_model(**kwargs)
return model_outputs.last_hidden_state

1045
src/ctx_to_lora/modeling_utils.py → src/ctx_to_lora/modeling/hypernet.py
View file

File diff suppressed because it is too large Load diff

57
src/ctx_to_lora/modeling_idefics2.py → src/ctx_to_lora/modeling/idefics2.py
View file

@ -1,4 +1,3 @@
# coding=utf-8
# Copyright 2024 the HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
@ -15,32 +14,22 @@
"""PyTorch Idefics2 model."""
import math
from dataclasses import dataclass
from typing import List, Optional, Tuple, Union
import torch
import torch.utils.checkpoint
from torch import nn
from torch.nn import CrossEntropyLoss
from transformers.activations import ACT2FN
from transformers.cache_utils import Cache, DynamicCache
from transformers.generation import GenerationMixin
from transformers.cache_utils import Cache
from transformers.configuration_utils import PretrainedConfig
from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask
from transformers.modeling_outputs import BaseModelOutput, ModelOutput
from transformers.modeling_utils import PreTrainedModel, ALL_ATTENTION_FUNCTIONS
from transformers.modeling_utils import PreTrainedModel
from transformers.models.idefics2.configuration_idefics2 import Idefics2Config
from transformers.utils import (
add_start_docstrings,
add_start_docstrings_to_model_forward,
is_flash_attn_2_available,
is_flash_attn_greater_or_equal_2_10,
logging,
replace_return_docstrings,
)
from transformers.models.auto import AutoModel
from transformers.configuration_utils import PretrainedConfig
from transformers.models.idefics2.configuration_idefics2 import Idefics2Config
if is_flash_attn_2_available():
from transformers.modeling_flash_attention_utils import _flash_attention_forward
@ -223,7 +212,7 @@ class Idefics2RMSNorm(nn.Module):
class Idefics2PerceiverAttention(nn.Module):
def __init__(self, config, layer_idx: Optional[int] = None) -> None:
def __init__(self, config, layer_idx: int | None = None) -> None:
"""Perceiver Cross-Attention Module --> let long-form inputs be `context`, resampled embeddings be `latents`"""
super().__init__()
self.config = config
@ -254,12 +243,12 @@ class Idefics2PerceiverAttention(nn.Module):
self,
latents: torch.Tensor,
context: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
past_key_value: Optional[Tuple[torch.Tensor]] = None,
attention_mask: torch.Tensor | None = None,
position_ids: torch.LongTensor | None = None,
past_key_value: tuple[torch.Tensor] | None = None,
output_attentions: bool = False,
use_cache: bool = False,
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
"""
Runs Perceiver Self-Attention, with special (context, latents) appended along the `seq` dimension!
@ -365,13 +354,13 @@ class Idefics2PerceiverFlashAttention2(Idefics2PerceiverAttention):
self,
latents: torch.Tensor,
context: torch.Tensor,
attention_mask: Optional[torch.LongTensor] = None,
position_ids: Optional[torch.LongTensor] = None,
past_key_value: Optional[Cache] = None,
attention_mask: torch.LongTensor | None = None,
position_ids: torch.LongTensor | None = None,
past_key_value: Cache | None = None,
output_attentions: bool = False,
use_cache: bool = False,
**kwargs,
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
bsz, q_len, _ = latents.size()
# kv_seq_len = q_len + context.size()[1]
kv_seq_len = context.size()[1]
@ -526,13 +515,13 @@ class Idefics2PerceiverLayer(nn.Module):
self,
latents: torch.Tensor,
context: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
past_key_value: Optional[Tuple[torch.Tensor]] = None,
output_attentions: Optional[bool] = False,
use_cache: Optional[bool] = False,
attention_mask: torch.Tensor | None = None,
position_ids: torch.LongTensor | None = None,
past_key_value: tuple[torch.Tensor] | None = None,
output_attentions: bool | None = False,
use_cache: bool | None = False,
**kwargs,
) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]:
"""
Args:
latents (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
@ -624,8 +613,8 @@ class Idefics2PerceiverResampler(Idefics2PreTrainedModel):
def forward(
self,
context: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
attention_mask: torch.Tensor | None = None,
position_ids: torch.LongTensor | None = None,
) -> torch.Tensor:
# seq embed -> bsz seq embed
if position_ids is None:
@ -732,8 +721,8 @@ class Idefics2Perceiver(Idefics2PreTrainedModel):
def forward(
self,
context: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
attention_mask: torch.Tensor | None = None,
position_ids: torch.LongTensor | None = None,
):
if position_ids is None:
bsz = context.shape[0]

9
src/ctx_to_lora/pooling.py
View file

@ -1,10 +1,8 @@
from enum import Enum
from typing import Optional
import torch
from jaxtyping import Float, Integer
from torch import Tensor
from torch.nn import functional as F
POOL_FN = Enum("POOL_FN", ["MEAN", "MAX", "LAST_TOKEN"])
@ -24,7 +22,7 @@ def get_pooling_fn(pooling_type: str):
def mean_pool(
features: Float[Tensor, "bs seq_len feature_dim"],
attn_mask: Optional[Integer[Tensor, "bs seq_len"]] = None,
attn_mask: Integer[Tensor, "bs seq_len"] | None = None,
) -> Float[Tensor, "bs 1 feature_dim"]:
if attn_mask is not None:
features = features.masked_fill(inv_bool_mask(attn_mask), 0)
@ -33,7 +31,7 @@ def mean_pool(
def max_pool(
features: Float[Tensor, "bs seq_len feature_dim"],
attn_mask: Optional[Integer[Tensor, "bs seq_len"]] = None,
attn_mask: Integer[Tensor, "bs seq_len"] | None = None,
) -> Float[Tensor, "bs 1 feature_dim"]:
if attn_mask is not None:
features = features.masked_fill(inv_bool_mask(attn_mask), -float("inf"))
@ -42,9 +40,8 @@ def max_pool(
def last_token_pool(
features: Float[Tensor, "bs seq_len feature_dim"],
attn_mask: Optional[Integer[Tensor, "bs seq_len"]] = None,
attn_mask: Integer[Tensor, "bs seq_len"] | None = None,
) -> Float[Tensor, "bs feature_dim"]:
left_padding = attn_mask[:, -1].sum() == attn_mask.shape[0]
if left_padding:
return features[:, -1]

5
src/ctx_to_lora/self_generate_data.py
View file

@ -1,5 +0,0 @@
# when generating dont forget to remove dup contexts
# merge QAs from dup contexts to construct
# a doc with QAs at the end
# maybe mix with some that doesn;t have QAs
# TODO: implement

99
src/ctx_to_lora/training_utils.py → src/ctx_to_lora/trainer.py
View file

@ -1,16 +1,11 @@
import gc
import logging
from enum import Enum
import torch
from transformers import Trainer
from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
from transformers.trainer import _is_peft_model
from transformers.trainer_utils import IntervalStrategy
from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
from ctx_to_lora.modeling_utils import ModulatedPretrainedModel
TRAINING_TASK = Enum("TRAINING_TASK", ["CAUSAL_LM", "COMPLETION"])
from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
logger = logging.getLogger()
@ -20,10 +15,12 @@ class ModulatedModelTrainer(Trainer):
self.gen_lora_l1_reg_coef = kwargs.pop("gen_lora_l1_reg_coef", 0.0)
super().__init__(*args, **kwargs)
def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
def compute_loss(
self, model, inputs, return_outputs=False, num_items_in_batch=None
):
"""
How the loss is computed by Trainer. By default, all models return the loss in the first element.
How the loss is computed by Trainer.
By default, all models return the loss in the first element.
Subclass and override for custom behavior.
"""
if (
@ -61,8 +58,11 @@ class ModulatedModelTrainer(Trainer):
else:
if isinstance(outputs, dict) and "loss" not in outputs:
raise ValueError(
"The model did not return a loss from the inputs, only the following keys: "
f"{','.join(outputs.keys())}. For reference, the inputs it received are {','.join(inputs.keys())}."
"The model did not return a loss from the inputs, "
"only the following keys: "
f"{','.join(outputs.keys())}. "
"For reference, the inputs it received are "
f"{','.join(inputs.keys())}."
)
# We don't use .loss here since the model may return tuples instead of ModelOutput.
loss = outputs["loss"] if isinstance(outputs, dict) else outputs[0]
@ -86,27 +86,13 @@ class ModulatedModelTrainer(Trainer):
return (loss, outputs) if return_outputs else loss
def clear_gpu():
gc.collect()
torch.cuda.empty_cache()
torch.cuda.reset_max_memory_allocated()
torch.cuda.reset_max_memory_cached()
def train_model(
model,
# tokenizer,
training_args,
train_dataset=None,
val_dataset=None,
# test_dataset=None,
train_collator=None,
# generation_collator=None,
compute_metrics=None,
train_sampler=None,
# preprocess_logits_for_metrics=None,
# max_new_tokens=2**13,
# gen_per_device_eval_batch_size=1,
):
checkpoint = None
if training_args.resume_from_checkpoint is not None:
@ -135,65 +121,6 @@ def train_model(
train_result = trainer.train(resume_from_checkpoint=checkpoint)
trainer.log_metrics("train", train_result.metrics)
trainer.save_model()
clear_gpu()
# metrics = trainer.evaluate(dict(**val_dataset, test=test_dataset))
# trainer.log_metrics("eval", metrics)
# trainer.save_metrics("eval", metrics)
# trainer.save_model()
# TODO: add benchmark eval?
# clear_gpu()
# ############## Evaluation
# # TODO: eval does not work when using with deepspeed
# # make a separate eval script
# # max_input_len=2**13 # for input truncation
# gen_kwargs = dict(do_sample=False, max_new_tokens=max_new_tokens)
# # pad_token_id=tokenizer.pad_token_id,
# # eos_token_id=?
# eval_trainer_args = {}
# # Copy only necessary attributes from training_args to eval_trainer_args
# seq2seq_training_args_fields = {f.name for f in fields(Seq2SeqTrainingArguments)}
# for attr, value in training_args.to_dict().items():
# if attr in seq2seq_training_args_fields:
# eval_trainer_args[attr] = value
# eval_trainer_args["eval_strategy"] = "no"
# eval_trainer_args["save_strategy"] = "no"
# eval_trainer_args["overwrite_output_dir"] = True
# eval_trainer_args["per_device_eval_batch_size"] = gen_per_device_eval_batch_size
# # NOTE: could also set kv_cache implementation here
# eval_trainer_args = Seq2SeqTrainingArguments(
# **eval_trainer_args,
# predict_with_generate=True,
# generation_config=GenerationConfig(**gen_kwargs),
# )
# # Seq2SeqTrainer is actually just the same as Trainer
# # (although it uses a different data collator, i.e., explicit prompt/answer separation)
# # it just allows `predict_with_generate`
# # allowing us to compute metrics on the generated outputs
# # no clue why they call this seq2seq...
# logger.info("=" * 80 + "\n" + "Evaluating model..." + "\n" + "=" * 80)
# model.eval()
# eval_trainer = Seq2SeqTrainer(
# model=model,
# args=eval_trainer_args,
# # TODO: use a different collator for test, e.g., more max_len truncation
# # w/ left padding?
# # removing label part from input_ids
# data_collator=generation_collator,
# )
# for split, ds in zip(["eval", "test"], [val_dataset, test_dataset]):
# if ds is None:
# continue
# eval_generation(eval_trainer, tokenizer, ds, split, gen_kwargs)
# clear_gpu()

27
src/ctx_to_lora/utils.py
View file

@ -1,12 +1,14 @@
import ast
import gc
import hashlib
import logging
import os
import random
import string
import time
import hashlib
from collections.abc import Iterable
from contextlib import contextmanager
from typing import Iterable, Optional
from enum import Enum
import torch
import yaml
@ -14,6 +16,9 @@ from peft import PeftConfig, PeftModel
from peft.tuners.tuners_utils import BaseTunerLayer, check_target_module_exists
from peft.utils import get_peft_model_state_dict
TRAINING_TASK = Enum("TRAINING_TASK", ["CAUSAL_LM", "COMPLETION"])
logger = logging.getLogger()
@ -74,7 +79,7 @@ def log_num_train_params(model):
)
def get_run_name(seed_str: Optional[str] = None):
def get_run_name(seed_str: str | None = None):
if not seed_str:
uuid = "".join(
[random.choice(string.ascii_letters + string.digits) for _ in range(8)]
@ -158,9 +163,8 @@ def save_yaml(data, path):
def get_peft_in_out_features(
model: PeftModel,
peft_config: Optional[PeftConfig] = None,
peft_config: PeftConfig | None = None,
) -> tuple[dict[str, int], dict[str, int]]:
if peft_config is None:
return None, None
in_features = dict()
@ -172,9 +176,9 @@ def get_peft_in_out_features(
continue
# support just Linear layer for now
# all modules should be a leave module that is Linear layer
assert isinstance(
module.base_layer, torch.nn.Linear
), "all modules should be a leave module that is Linear layer"
assert isinstance(module.base_layer, torch.nn.Linear), (
"all modules should be a leave module that is Linear layer"
)
# this should always pass
name = module_name.split(".")[-1]
@ -233,3 +237,10 @@ def get_lora_module_names(
module_names[target_module][layer_idx].append(k)
break
return module_names
def clear_gpu():
gc.collect()
torch.cuda.empty_cache()
torch.cuda.reset_max_memory_allocated()
torch.cuda.reset_max_memory_cached()

4257
uv.lock generated Normal file
View file

File diff suppressed because it is too large Load diff

22
watcher.py
View file

@ -2,23 +2,19 @@
# not using watchdog because there are too many saved files
# but we want to just watch when these files are created
# */checkpoints/it_*/hypermod.pt (for HyperLoRA)
import itertools
import shutil
import time
import os
import argparse
import gc
import itertools
import os
import time
from glob import glob
import numpy as np
import pandas as pd
import wandb
import yaml
import torch
import yaml
import wandb
from eval import evaluate
CP_PATTERN = "train_outputs/runs/*/checkpoint*/pytorch_model.bin"
@ -55,11 +51,13 @@ class Watcher:
def load_state(self):
if not os.path.exists("watcher_state.yaml"):
return
with open("watcher_state.yaml", "r") as f:
with open("watcher_state.yaml") as f:
state = yaml.safe_load(f)
self.last_files = state["last_files"]
# TODO: run eval on "validation" split during training
# using function from `src/ctx-to-lora/eval.py`
if __name__ == "__main__":
os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true"
os.environ["FLASH_ATTENTION_DETERMINISTIC"] = "1"
@ -94,9 +92,7 @@ if __name__ == "__main__":
run_dir = file.split("/checkpoint")[0]
cur_it = int(file.split("checkpoint-")[1].split("/")[0])
checkpoint_dir = os.path.dirname(file)
args = argparse.Namespace(
**yaml.unsafe_load(open(f"{run_dir}/args.yaml", "r"))
)
args = argparse.Namespace(**yaml.unsafe_load(open(f"{run_dir}/args.yaml")))
args.output_dir = f"{run_dir}/eval-results-{cur_it}"
args.logging_dir = f"{run_dir}/eval-results-{cur_it}"
args.run_name = run_dir.split("/")[-1]

35
webui/app.py
View file

@ -1,14 +1,13 @@
import json
import os
import re
from glob import glob
import torch
import yaml
from flask import Flask, abort, jsonify, render_template, request
from transformers import pipeline, AutoTokenizer
from transformers import pipeline
from ctx_to_lora.data_utils import tokenize_ctx_text
from ctx_to_lora.data.processing import tokenize_ctx_text
from ctx_to_lora.model_loading import get_tokenizer
app = Flask(__name__)
@ -132,11 +131,9 @@ def get_generated_text_data(run_path: str) -> dict:
generated_data = {}
# Find all *_generated_text.jsonl files in the current directory, excluding _no_context_ variants
files = sorted(
[
f
for f in glob(os.path.join(run_path, "*_generated_text.jsonl"))
if "_no_context_generated_text.jsonl" not in os.path.basename(f)
]
f
for f in glob(os.path.join(run_path, "*_generated_text.jsonl"))
if "_no_context_generated_text.jsonl" not in os.path.basename(f)
)
print(f"Processing generated text files from {run_path}: {files}")
@ -159,11 +156,9 @@ def get_generated_text_data(run_path: str) -> dict:
]:
subdir_name = os.path.basename(subdir_path)
subdir_files = sorted(
[
f
for f in glob(os.path.join(subdir_path, "*_generated_text.jsonl"))
if "_no_context_generated_text.jsonl" not in os.path.basename(f)
]
f
for f in glob(os.path.join(subdir_path, "*_generated_text.jsonl"))
if "_no_context_generated_text.jsonl" not in os.path.basename(f)
)
print(
f"Processing generated text files from subdirectory {subdir_path}: {subdir_files}"
@ -259,7 +254,7 @@ def get_available_checkpoints(run):
return []
checkpoints = sorted(
[d for d in glob(os.path.join(logdir, "checkpoint-*")) if os.path.isdir(d)],
(d for d in glob(os.path.join(logdir, "checkpoint-*")) if os.path.isdir(d)),
key=lambda d: int(d.split("-")[-1]),
)
@ -282,7 +277,7 @@ def load_custom_chat_template(tokenizer, model_name):
if "gemma" in model_name.lower():
template_path = "chat_templates/google/gemma-2-2b-it.jinja"
if os.path.exists(template_path):
with open(template_path, "r") as f:
with open(template_path) as f:
template_content = f.read()
tokenizer.chat_template = template_content
print(f"Loaded custom chat template from {template_path}")
@ -292,7 +287,7 @@ def load_custom_chat_template(tokenizer, model_name):
elif "meta-llama" in model_name.lower() or "llama" in model_name.lower():
template_path = "chat_templates/meta-llama/llama-2-7b-chat.jinja"
if os.path.exists(template_path):
with open(template_path, "r") as f:
with open(template_path) as f:
template_content = f.read()
tokenizer.chat_template = template_content
print(f"Loaded custom chat template from {template_path}")
@ -478,7 +473,7 @@ def load_checkpoint():
# Unload the regular chat model to avoid confusion
chat_generator = None
from ctx_to_lora.modeling_utils import ModulatedPretrainedModel
from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
run = request.form["run"]
checkpoint = request.form["checkpoint"]
@ -592,8 +587,6 @@ def chat():
# Process context and generate LoRA for the response
ctx_tokenizer = None
with torch.inference_mode(), torch.amp.autocast(str(device)):
from transformers import AutoTokenizer
# Get the contexts and tokenize them
contexts = request.form.getlist("contexts[]")
if not contexts:
@ -603,7 +596,9 @@ def chat():
if len(contexts) == 1 and not contexts[0].strip():
contexts = [""]
print(f"Processing {len(contexts)} contexts for response generation")
print(
f"Processing {len(contexts)} contexts for response generation"
)
print(f"Contexts: {contexts}")
ctx_encoder_model_name_or_path = (
modulated_model.ctx_encoder_args.ctx_encoder_model_name_or_path