doc-to-lora/data/generate_negative_natural_questions.py
Rujikorn Charakorn 891c0bd256
Refactor_and_improve_data (#7)
* faster slice

* add facts + ctx_qa

* new configs

* new scripts

* intx_sft.py to train.py

* add kaggle for downloading facts

* max_new_tokens cli for eval

* generate negative_nq

* scripts + configs

* default vals

* max_val_samples_per_ds=500

* more efficient layer-to-layer ctx encoder

* use_per_ctx_average_loss

* faster processing

* small exp distill

* scripts

* more robust watcher

* per-module l1_norm avg

* per-ctx average loss

* clear_gpu
2025-08-04 20:52:37 +09:00

44 lines
1.6 KiB
Python

import os
from datasets import Dataset, load_dataset
from tqdm import tqdm
if __name__ == "__main__":
ds = load_dataset("google-research-datasets/natural_questions")
ds = ds.shuffle(seed=42)
for split in ["validation"]:
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")