From b1f92663a700a8c2e18a549c45b698b85756207a Mon Sep 17 00:00:00 2001 From: 51616 Date: Fri, 30 May 2025 03:11:03 +0000 Subject: [PATCH] data gen scripts --- .../generate_data.py => generate_fav.py} | 0 ...te_data_big.py => generate_fav_num_big.py} | 0 data/generate_fw_edu_qa_vllm.py | 235 ++++++++++++++++++ data/generate_negative_natural_questions.py | 44 ++++ data/self_generate_qa.py | 136 ++++++++++ 5 files changed, 415 insertions(+) rename data/{raw_datasets/generate_data.py => generate_fav.py} (100%) rename data/{raw_datasets/generate_data_big.py => generate_fav_num_big.py} (100%) create mode 100644 data/generate_fw_edu_qa_vllm.py create mode 100644 data/generate_negative_natural_questions.py create mode 100644 data/self_generate_qa.py diff --git a/data/raw_datasets/generate_data.py b/data/generate_fav.py similarity index 100% rename from data/raw_datasets/generate_data.py rename to data/generate_fav.py diff --git a/data/raw_datasets/generate_data_big.py b/data/generate_fav_num_big.py similarity index 100% rename from data/raw_datasets/generate_data_big.py rename to data/generate_fav_num_big.py diff --git a/data/generate_fw_edu_qa_vllm.py b/data/generate_fw_edu_qa_vllm.py new file mode 100644 index 0000000..451dddc --- /dev/null +++ b/data/generate_fw_edu_qa_vllm.py @@ -0,0 +1,235 @@ +import os +import random +import re +import sys +from glob import glob + +import pandas as pd +from datasets import Dataset, load_dataset +from vllm import LLM, SamplingParams + +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() diff --git a/data/generate_negative_natural_questions.py b/data/generate_negative_natural_questions.py new file mode 100644 index 0000000..660aac8 --- /dev/null +++ b/data/generate_negative_natural_questions.py @@ -0,0 +1,44 @@ +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 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") diff --git a/data/self_generate_qa.py b/data/self_generate_qa.py new file mode 100644 index 0000000..46ed082 --- /dev/null +++ b/data/self_generate_qa.py @@ -0,0 +1,136 @@ +import os +import random +import sys +from glob import glob + +import pandas as pd +from datasets import Dataset, load_dataset +from vllm import LLM, SamplingParams + +SYSTEM_TEMPLATE = ( + "### System Instructions ###\n" + "You are a creative and helpful assistant.\n" + "**DO NOT** halucinate or make up information." +) + +PROMPT_TEMPLATE = "### Context ###\n{context}\n\n\n### Question ###\n{question}" + +MODEL_CTX_LEN = { + "google/gemma-2-27b-it": 8192, + "google/gemma-2-2b-it": 8192, + "google/gemma-2-9b-it": 8192, +} + + +def get_prompt(context, q): + prompt = PROMPT_TEMPLATE.format(context=context, question=q) + return prompt + + +if __name__ == "__main__": + vllm_model = os.environ.get("vllm_model") # "google/gemma-2-27b-it" + debug = bool(os.environ.get("debug", False)) + print(f"Using model: {vllm_model}") + llm_kwargs = dict( + model=vllm_model, + dtype="bfloat16", + enable_prefix_caching=True, + enable_chunked_prefill=True, + max_model_len=MODEL_CTX_LEN.get(vllm_model), + ) + + if vllm_model == "mistralai/Mistral-Small-3.1-24B-Instruct-2503": + mm_kwargs = {"image": 0} # if "gemma-3" in vllm_model else {} + 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_kwargs["max_model_len"] = 2**14 + llm_kwargs["limit_mm_per_prompt"] = mm_kwargs + + print(f"{llm_kwargs=}") + llm = LLM(**llm_kwargs) + tokenizer = llm.get_tokenizer() + shard_pattern = sys.argv[1] + + paths = glob(f"./data/raw_datasets/fw_qa_3/{shard_pattern}.parquet") + + for path in paths: + if debug: + ds = load_dataset( + "parquet", + data_files=path, + split="train[:10]", + ) + else: + ds = load_dataset( + "parquet", + data_files=path, + split="train", + # streaming=True, + ) + ctxs = [] + questions = [] + for sample in iter(ds): + ctxs.append(sample["context"]) + questions.append(sample["prompt"]) + if "gemma" in vllm_model: + # gemma models do not support system messages + messages = [ + [ + { + "role": "user", + "content": SYSTEM_TEMPLATE + "\n\n\n" + get_prompt(ctx, q), + }, + ] + for ctx, q in zip(ctxs, questions) + ] + else: + messages = [ + [ + {"role": "system", "content": SYSTEM_TEMPLATE}, + {"role": "user", "content": get_prompt(ctx, q)}, + ] + for ctx, q in zip(ctxs, questions) + ] + + 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 1.0 + ), + ), + ) + samples = [] + for ctx, q, completion in zip(ctxs, questions, completions): + response = completion.outputs[0].text + samples.append( + { + "context": ctx, + "prompt": q, + "response": response, + } + ) + if debug: + print(f"{ctx=}") + print(f"{q=}") + print(f"{response=}") + print("=" * 80) + + print(f"Generated {len(samples)} samples") + random.shuffle(samples) + df = pd.DataFrame(samples) + ds = Dataset.from_pandas(df) + + shard_name = path.split("/")[-1].split(".")[0] + ds.to_parquet(f"data/raw_datasets/fw_qa_3/{vllm_model}/{shard_name}.parquet") + print(f"Saved to data/raw_datasets/fw_qa_3/{vllm_model}/{shard_name}.parquet")