Initial commit

scripts/data_process/nq.py

@@ -0,0 +1,100 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Preprocess the nq dataset to parquet format
+"""
+
+import re
+import os
+import datasets
+
+from verl.utils.hdfs_io import copy, makedirs
+import argparse
+
+
+def make_prefix(dp, template_type):
+    question = dp['question']
+
+    # NOTE: also need to change reward_score/countdown.py
+    if template_type == 'base':
+        """This works for any base model"""
+        prefix = f"""Answer the given question. \
+You should first have a reasoning process in mind and then provides the answer. \
+Show your reasoning in <think> </think> tags and return the final answer in <answer> </answer> tags, for example <answer> Beijing </answer>. \
+Question: {question}\n"""
+    else:
+        raise NotImplementedError    
+    return prefix
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--local_dir', default='./data/nq')
+    parser.add_argument('--hdfs_dir', default=None)
+    parser.add_argument('--template_type', type=str, default='base')
+
+    args = parser.parse_args()
+
+    data_source = 'nq'
+
+    dataset = datasets.load_dataset('RUC-NLPIR/FlashRAG_datasets', 'nq')
+
+    train_dataset = dataset['train']
+    test_dataset = dataset['test']
+
+    # add a row to each data item that represents a unique id
+    def make_map_fn(split):
+
+        def process_fn(example, idx):
+            example['question'] = example['question'].strip()
+            if example['question'][-1] != '?':
+                example['question'] += '?'
+            question = make_prefix(example, template_type=args.template_type)
+            solution = {
+                "target": example['golden_answers'],
+            }
+
+            data = {
+                "data_source": data_source,
+                "prompt": [{
+                    "role": "user",
+                    "content": question,
+                }],
+                "ability": "fact-reasoning",
+                "reward_model": {
+                    "style": "rule",
+                    "ground_truth": solution
+                },
+                "extra_info": {
+                    'split': split,
+                    'index': idx,
+                }
+            }
+            return data
+
+        return process_fn
+
+    train_dataset = train_dataset.map(function=make_map_fn('train'), with_indices=True)
+    test_dataset = test_dataset.map(function=make_map_fn('test'), with_indices=True)
+
+    local_dir = args.local_dir
+    hdfs_dir = args.hdfs_dir
+
+    train_dataset.to_parquet(os.path.join(local_dir, 'train.parquet'))
+    test_dataset.to_parquet(os.path.join(local_dir, 'test.parquet'))
+
+    if hdfs_dir is not None:
+        makedirs(hdfs_dir)
+
+        copy(src=local_dir, dst=hdfs_dir)

scripts/data_process/nq_rag.py

@@ -0,0 +1,141 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Preprocess the nq dataset to parquet format
+"""
+
+import re
+import os
+import json
+import datasets
+
+from verl.utils.hdfs_io import copy, makedirs
+import argparse
+
+
+def make_prefix(dp, template_type):
+    question = dp['question']
+    context = dp['context']
+
+    # NOTE: also need to change reward_score/countdown.py
+    if template_type == 'base':
+        """This works for any base model"""
+        prefix = f"""Answer the given question with some potentially useful context. \
+You should analyze the question carefully, evaluate the given context (which may or may not be useful), and then generate an accurate and well-reasoned response. \
+You should first have a reasoning process in mind and then provides the answer. \
+Show your reasoning in <think> </think> tags and return the final answer in <answer> </answer> tags, for example <answer> Beijing </answer>. \
+Question: {question} Context: {context} \n"""
+    else:
+        raise NotImplementedError    
+    return prefix
+
+
+def format_reference(retrieval_result):
+    format_reference = ''
+    for idx, doc_item in enumerate(retrieval_result):
+        content = doc_item['contents']
+        title = content.split("\n")[0]
+        text = "\n".join(content.split("\n")[1:])
+        format_reference += f"Doc {idx+1}(Title: {title}) {text}\n"
+
+    return format_reference
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--local_dir', default='./data/nq_rag')
+    parser.add_argument('--hdfs_dir', default=None)
+    parser.add_argument('--template_type', type=str, default='base')
+    parser.add_argument('--topk', type=int, default=3)
+    parser.add_argument('--corpus_path', type=str, default='/home/peterjin/mnt/data/retrieval-corpus/wiki-18.jsonl')
+    parser.add_argument('--train_retrieval_cache', type=str, default='/home/peterjin/rag_retrieval_cache/nq/e5_train_retrieval_cache_2048.json')
+    parser.add_argument('--test_retrieval_cache', type=str, default='/home/peterjin/rag_retrieval_cache/nq/e5_test_retrieval_cache_10000.json')
+
+    args = parser.parse_args()
+
+    data_source = 'nq'
+
+    dataset = datasets.load_dataset('RUC-NLPIR/FlashRAG_datasets', 'nq')
+
+    train_dataset = dataset['train']
+    test_dataset = dataset['test']
+
+    # read retrieval cache
+    print('reading retrieval cache...')
+    retrieval_cache = json.load(open(args.train_retrieval_cache))
+    # test_retrieval_cache = json.load(open(args.test_retrieval_cache))
+    retrieval_cache.update(json.load(open(args.test_retrieval_cache)))
+
+    # read corpus
+    print('reading corpus...')
+    corpus = {}
+    with open(args.corpus_path) as f:
+        readin = f.readlines()
+        for line in readin:
+            tmp = json.loads(line)
+            corpus[tmp['id']] = tmp
+
+    # add a column for the retrieval context
+    def add_context(example):
+        example['context'] = format_reference([corpus[docs["id"]] for docs in retrieval_cache[example['question']][:args.topk]])
+        return example
+    
+    train_dataset = train_dataset.map(function=add_context)
+    test_dataset = test_dataset.map(function=add_context)
+
+    # add a row to each data item that represents a unique id
+    def make_map_fn(split):
+
+        def process_fn(example, idx):
+            example['question'] = example['question'].strip()
+            if example['question'][-1] != '?':
+                example['question'] += '?'
+            question = make_prefix(example, template_type=args.template_type)
+            solution = {
+                "target": example['golden_answers'],
+            }
+
+            data = {
+                "data_source": data_source,
+                "prompt": [{
+                    "role": "user",
+                    "content": question,
+                }],
+                "ability": "fact-reasoning",
+                "reward_model": {
+                    "style": "rule",
+                    "ground_truth": solution
+                },
+                "extra_info": {
+                    'split': split,
+                    'index': idx,
+                }
+            }
+            return data
+
+        return process_fn
+
+    train_dataset = train_dataset.map(function=make_map_fn('train'), with_indices=True)
+    test_dataset = test_dataset.map(function=make_map_fn('test'), with_indices=True)
+
+    local_dir = args.local_dir
+    hdfs_dir = args.hdfs_dir
+
+    train_dataset.to_parquet(os.path.join(local_dir, 'train.parquet'))
+    test_dataset.to_parquet(os.path.join(local_dir, 'test.parquet'))
+
+    if hdfs_dir is not None:
+        makedirs(hdfs_dir)
+
+        copy(src=local_dir, dst=hdfs_dir)

scripts/data_process/nq_search.py

@@ -0,0 +1,101 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Preprocess the nq dataset to parquet format
+"""
+
+import re
+import os
+import datasets
+
+from verl.utils.hdfs_io import copy, makedirs
+import argparse
+
+
+def make_prefix(dp, template_type):
+    question = dp['question']
+
+    # NOTE: also need to change reward_score/countdown.py
+    if template_type == 'base':
+        """This works for any base model"""
+        prefix = f"""Answer the given question. \
+You must conduct reasoning inside <think> and </think> first every time you get new information. \
+After reasoning, if you find you lack some knowledge, you can call a search engine by <search> query </search> and it will return the top searched results between <information> and </information>. \
+You can search as many times as your want. \
+If you find no further external knowledge needed, you can directly provide the answer inside <answer> and </answer>, without detailed illustrations. For example, <answer> Beijing </answer>. Question: {question}\n"""
+    else:
+        raise NotImplementedError
+    return prefix
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--local_dir', default='./data/nq_search')
+    parser.add_argument('--hdfs_dir', default=None)
+    parser.add_argument('--template_type', type=str, default='base')
+
+    args = parser.parse_args()
+
+    data_source = 'nq'
+
+    dataset = datasets.load_dataset('RUC-NLPIR/FlashRAG_datasets', 'nq')
+
+    train_dataset = dataset['train']
+    test_dataset = dataset['test']
+
+    # add a row to each data item that represents a unique id
+    def make_map_fn(split):
+
+        def process_fn(example, idx):
+            example['question'] = example['question'].strip()
+            if example['question'][-1] != '?':
+                example['question'] += '?'
+            question = make_prefix(example, template_type=args.template_type)
+            solution = {
+                "target": example['golden_answers'],
+            }
+
+            data = {
+                "data_source": data_source,
+                "prompt": [{
+                    "role": "user",
+                    "content": question,
+                }],
+                "ability": "fact-reasoning",
+                "reward_model": {
+                    "style": "rule",
+                    "ground_truth": solution
+                },
+                "extra_info": {
+                    'split': split,
+                    'index': idx,
+                }
+            }
+            return data
+
+        return process_fn
+
+    train_dataset = train_dataset.map(function=make_map_fn('train'), with_indices=True)
+    test_dataset = test_dataset.map(function=make_map_fn('test'), with_indices=True)
+
+    local_dir = args.local_dir
+    hdfs_dir = args.hdfs_dir
+
+    train_dataset.to_parquet(os.path.join(local_dir, 'train.parquet'))
+    test_dataset.to_parquet(os.path.join(local_dir, 'test.parquet'))
+
+    if hdfs_dir is not None:
+        makedirs(hdfs_dir)
+
+        copy(src=local_dir, dst=hdfs_dir)

scripts/download.py

@@ -0,0 +1,25 @@
+import argparse
+from huggingface_hub import hf_hub_download
+
+parser = argparse.ArgumentParser(description="Download files from a Hugging Face dataset repository.")
+parser.add_argument("--repo_id", type=str, default="PeterJinGo/wiki-18-e5-index", help="Hugging Face repository ID")
+parser.add_argument("--save_path", type=str, required=True, help="Local directory to save files")
+    
+args = parser.parse_args()
+
+repo_id = "PeterJinGo/wiki-18-e5-index"
+for file in ["part_aa", "part_ab"]:
+    hf_hub_download(
+        repo_id=repo_id,
+        filename=file,  # e.g., "e5_Flat.index"
+        repo_type="dataset",
+        local_dir=args.save_path,
+    )
+
+repo_id = "PeterJinGo/wiki-18-corpus"
+hf_hub_download(
+        repo_id=repo_id,
+        filename="wiki-18.jsonl.gz",
+        repo_type="dataset",
+        local_dir=args.save_path,
+)

scripts/download.sh

@@ -0,0 +1,6 @@
+
+save_path=/home/peterjin/debug_cache
+
+python download.py --savepath $savepath
+
+cat $save_path/part_* > e5_Flat.index

scripts/upload.py

@@ -0,0 +1,12 @@
+import os
+from huggingface_hub import upload_file
+
+repo_id = "PeterJinGo/wiki-18-e5-index"
+path = "/home/peterjin/mnt/index/wiki-18"
+for file in ["part_aa", "part_ab"]:
+    upload_file(
+        path_or_fileobj=os.path.join(path, file),  # File path
+        path_in_repo=file,  # Destination filename in the repo
+        repo_id=repo_id,  # Your dataset repo ID
+        repo_type="dataset"
+    )

scripts/upload.sh

@@ -0,0 +1,6 @@
+
+index=/home/peterjin/mnt/index/wiki-18/e5_Flat.index
+
+split -b 40G $index part_
+
+python upload.py