chore: import upstream snapshot with attribution

research/BGE_Coder/evaluation/coderag_eval/test/create/code_search_net.py

@@ -0,0 +1,52 @@
+import argparse
+import datasets
+import os
+from tqdm import tqdm
+import random
+from create.utils import save_tsv_dict, save_file_jsonl, load_jsonlines
+
+def document2code(data, split="train"):
+    data = data[split]
+    code_search_net_data_queries = []
+    code_search_net_data_docs = []
+    code_search_net_data_qrels = []
+
+    for item in tqdm(data):
+        doc = item["func_documentation_string"]
+        code = item["func_code_string"]
+        doc_id = "{repository_name}_{func_path_in_repository}_{func_name}_doc".format_map(item)
+        code_id = "{repository_name}_{func_path_in_repository}_{func_name}_code".format_map(item)
+        code_search_net_data_queries.append({"_id": doc_id, "text": doc, "metadata": {}})
+        code_search_net_data_docs.append({"_id": code_id, "title": item["func_name"], "text": code, "metadata": {}})
+        code_search_net_data_qrels.append({"query-id": doc_id, "corpus-id": code_id, "score": 1})
+
+    return code_search_net_data_queries, code_search_net_data_docs, code_search_net_data_qrels
+
+def main():
+#### /print debug information to stdout
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--language", type=str, default="python", help="codesearch net language")
+    parser.add_argument("--output_dir", type=str, default="datasets")
+
+    args = parser.parse_args()
+    dataset = datasets.load_dataset("code_search_net", args.language)
+
+    path = os.path.join(args.output_dir, "code_search_net_{}".format(args.language))
+    os.makedirs(path)
+    os.makedirs(os.path.join(path, "qrels"))
+
+    docs = []
+    queries = []
+    for split in ["train", "validation", "test"]:
+        queries_split, docs_split, qrels_split = document2code(dataset, split)
+        docs += docs_split
+        queries += queries_split
+
+        save_tsv_dict(qrels_split, os.path.join(path, "qrels", "{}.tsv".format(split)), ["query-id", "corpus-id", "score"])
+
+    save_file_jsonl(queries, os.path.join(path, "queries.jsonl"))
+    save_file_jsonl(docs, os.path.join(path, "corpus.jsonl"))
+
+if __name__ == "__main__":
+    main()

research/BGE_Coder/evaluation/coderag_eval/test/create/ds1000.py

@@ -0,0 +1,123 @@
+import io
+import os
+import fcntl
+import pathlib
+import zipfile
+import argparse
+import requests
+import warnings
+import itertools
+from tqdm import tqdm
+from datasets import load_dataset
+from create.utils import save_tsv_dict, save_file_jsonl
+
+
+# Load dataset
+def download_source(source_dir):
+    src = source_dir / "ds1000.py"
+    url = "https://github.com/HKUNLP/DS-1000/blob/49c1c543ada8b58138181333cdc62e613204efcf/ds1000.py?raw=true"
+    lock = src.with_suffix(".lock")
+    with open(lock, "w") as f_lock:
+        fcntl.flock(f_lock, fcntl.LOCK_EX)
+        if not src.exists():
+            warnings.warn(f"DS-1000 source is being saved to {src}.")
+            print("Downloading source code...")
+            r = requests.get(url, stream=True)
+            with open(src, "wb") as f_src:
+                f_src.write(r.content)
+            open(src.parent / "__init__.py", "w").close()
+            print("Done.")
+            fcntl.flock(f_lock, fcntl.LOCK_UN)
+
+def download_dataset(source_dir):
+    path = source_dir / "ds1000_data"
+    url = "https://github.com/HKUNLP/DS-1000/blob/49c1c543ada8b58138181333cdc62e613204efcf/ds1000_data.zip?raw=true"
+    lock = path.with_suffix(".lock")
+    with open(lock, "w") as f_lock:
+        fcntl.flock(f_lock, fcntl.LOCK_EX)
+        if not path.exists():
+            warnings.warn(f"DS-1000 data is being saved to {path}.")
+            print("Downloading dataset...")
+            r = requests.get(url, stream=True)
+            z = zipfile.ZipFile(io.BytesIO(r.content))
+            z.extractall(source_dir)
+            print("Done.")
+        fcntl.flock(f_lock, fcntl.LOCK_UN)
+
+def get_dataset(source_dir, mode: str = "Completion", key: str = "All"):
+    """Returns dataset for the task or an iterable of any object, that get_prompt can handle"""
+    from ds.ds1000 import DS1000Dataset
+
+    data = DS1000Dataset(source_dir / "ds1000_data", mode=mode).data
+    if key == "All":
+        if mode == "Insertion":
+            warnings.warn(
+                "Insertion not supported for Matplotlib. Only running others."
+            )
+            data = {k: v for k, v in data.items() if k != "Matplotlib"}
+        dataset = list(itertools.chain(*data.values()))
+    else:
+        dataset = data[key]
+    return dataset
+
+
+# Collect queries, docs, and relations
+def document2code(data: list):
+    queries, docs, qrels = [], [], []
+
+    # collect doc corpus
+    code_docs = load_dataset("neulab/docprompting-conala", "docs")["train"]
+    for i in range(len(code_docs)):
+        docs.append({
+            "_id": str(i),
+            "title": code_docs[i]["doc_id"],
+            "text": code_docs[i]["doc_content"],
+            "metadata": {}
+        })
+    
+    # load canonical docs
+    ds1000 = load_dataset("json", data_files={"test": args.canonical_file})["test"]
+    for idx,item in enumerate(tqdm(data)):
+        example = item.data
+        query = example["prompt"]
+        query_id = f"{example['lib']}_{example['perturbation_origin_id']}"
+        queries.append({"_id": query_id, "text": query, "metadata": {}})
+
+        doc_ids = [doc["title"] for doc in ds1000[idx]["docs"]]
+        for doc_id in doc_ids:
+            corpus_id = code_docs["doc_id"].index(doc_id)
+            corpus_id = str(corpus_id)
+            qrels.append({"query-id": query_id, "corpus-id": corpus_id, "score": 1})
+    
+    return queries, docs, qrels
+
+
+def main():
+    args.source_dir = pathlib.Path(__file__).parent.parent / args.source_dir
+    os.makedirs(args.source_dir, exist_ok=True)
+    download_source(args.source_dir)
+    download_dataset(args.source_dir)
+    dataset = get_dataset(args.source_dir, mode=args.mode, key=args.key)
+
+    path = os.path.join(args.output_dir, f"ds1000_{args.key.lower()}_{args.mode.lower()}")
+    os.makedirs(path, exist_ok=True)
+    os.makedirs(os.path.join(path, "qrels"), exist_ok=True)
+
+    queries, docs, qrels = document2code(dataset)
+    save_tsv_dict(qrels, os.path.join(path, "qrels", "test.tsv"), ["query-id", "corpus-id", "score"])
+    
+    save_file_jsonl(queries, os.path.join(path, "queries.jsonl"))
+    save_file_jsonl(docs, os.path.join(path, "corpus.jsonl"))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--source_dir", type=str, default="ds")
+    parser.add_argument("--output_dir", type=str, default="datasets")
+    parser.add_argument("--mode", type=str, default="Completion", choices=["Completion", "Insertion"])
+    parser.add_argument("--key", type=str, default="All", 
+                        choices=["All", "Numpy", "Pandas", "Scipy", "Matplotlib", "Sklearn", "Tensorflow", "Pytorch"])
+    parser.add_argument("--canonical_file", type=str, default="datasets/canonical/ds1000_docs.json")
+    args = parser.parse_args()
+
+    main()

research/BGE_Coder/evaluation/coderag_eval/test/create/general_programming.py

@@ -0,0 +1,70 @@
+"""Aggregate all code-generation datasets."""
+
+import os
+import json
+import datasets
+import argparse
+from create.utils import save_tsv_dict
+from create.humaneval import document2code as d2c_humaneval
+from create.mbpp import document2code as d2c_mbpp
+
+D2C_FUNC_DICT = {
+    "humaneval": d2c_humaneval,
+    "mbpp": d2c_mbpp,
+}
+SPLIT_DICT = {
+    "humaneval": ["test"], 
+    "mbpp": ["train", "test", "validation", "prompt"],
+}
+HF_NAME_DICT = {
+    "humaneval": "openai_humaneval",
+    "mbpp": "mbpp",
+}
+
+
+def save_file_jsonl(data, path):
+    with open(path,'w') as fw:
+        for item in data:
+            fw.write(json.dumps(item) + '\n')
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset_names", type=str, nargs='+', default=["humaneval", "mbpp"])
+    parser.add_argument("--output_dir", type=str, default="datasets")
+    parser.add_argument("--output_name", type=str, default="general-programming")
+    args = parser.parse_args()
+
+    path = os.path.join(args.output_dir, args.output_name)
+    os.makedirs(path)
+    os.makedirs(os.path.join(path, "qrels"), exist_ok=True)
+
+    split_dict = {}
+    for dataset_name in args.dataset_names:
+        for split in SPLIT_DICT[dataset_name]:
+            if split not in split_dict:
+                split_dict[split] = []
+            split_dict[split].append(dataset_name)
+    
+    dataset_dict = {
+        dataset_name: datasets.load_dataset(HF_NAME_DICT[dataset_name])
+        for dataset_name in args.dataset_names
+    }
+    docs, queries = [], []
+    for split, ds_names in split_dict.items():
+        for ds in ds_names:
+            dataset = dataset_dict[ds]
+
+            queries_split, docs_split, qrels_split = D2C_FUNC_DICT[ds](dataset, split)
+            docs += docs_split
+            queries += queries_split
+
+        qrels_path = os.path.join(path, "qrels", f"{split}.tsv")
+        save_tsv_dict(qrels_split, qrels_path, ["query-id", "corpus-id", "score"])
+    
+    save_file_jsonl(queries, os.path.join(path, "queries.jsonl"))
+    save_file_jsonl(docs, os.path.join(path, "corpus.jsonl"))
+
+
+if __name__ == "__main__":
+    main()

research/BGE_Coder/evaluation/coderag_eval/test/create/humaneval.py

@@ -0,0 +1,58 @@
+import os
+import argparse
+import datasets
+from tqdm import tqdm
+from create.utils import save_tsv_dict, save_file_jsonl
+
+
+def document2code(data, split="test"):
+    data = data[split]
+    queries, docs, qrels = [], [], []
+
+    for item in tqdm(data):
+        doc = item["prompt"]
+        code = item["prompt"] + '\n' + item["canonical_solution"]
+        doc_id = "{task_id}_doc".format_map(item)
+        code_id = "{task_id}_code".format_map(item)
+
+        queries.append({"_id": doc_id, "text": doc, "metadata": {}})
+        docs.append({"_id": code_id, "title": item["entry_point"], "text": code, "metadata": {}})
+        qrels.append({"query-id": doc_id, "corpus-id": code_id, "score": 1})
+    
+    return queries, docs, qrels
+
+
+def main():
+    dataset = datasets.load_dataset(args.dataset_name)
+
+    path = os.path.join(args.output_dir, args.output_name)
+    os.makedirs(path, exist_ok=True)
+    os.makedirs(os.path.join(path, "qrels"), exist_ok=True)
+
+    queries, docs, qrels = document2code(dataset, split="test")
+    save_file_jsonl(queries, os.path.join(path, "queries.jsonl"))
+    save_file_jsonl(docs, os.path.join(path, "corpus.jsonl"))
+    qrels_path = os.path.join(path, "qrels", "test.tsv")
+    save_tsv_dict(qrels, qrels_path, ["query-id", "corpus-id", "score"])
+
+    # create canonical file if not existent yet
+    if not os.path.exists(args.canonical_file):
+        canonical_solutions = []
+        for doc in docs:
+            canonical_solutions.append([{
+                "text": doc["text"], "title": doc["title"]
+            }])
+        canonical_dataset = dataset["test"].add_column("docs", canonical_solutions)
+        canonical_dataset.to_json(args.canonical_file)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset_name", type=str, default="openai_humaneval")
+    parser.add_argument("--output_name", type=str, default="humaneval")
+    parser.add_argument("--canonical_file", type=str, 
+                        default="datasets/canonical/humaneval_solutions.json")
+    parser.add_argument("--output_dir", type=str, default="datasets")
+    args = parser.parse_args()
+
+    main()

research/BGE_Coder/evaluation/coderag_eval/test/create/live_code_bench.py

@@ -0,0 +1,53 @@
+import os
+import argparse
+import datasets
+from tqdm import tqdm
+from datasets import load_dataset
+from create.utils import save_tsv_dict, save_file_jsonl
+
+
+def get_queries(data, split="test") -> list[dict]:
+    queries = [{
+        "_id": item["question_id"] + '__' + item["contest_id"], 
+        "text": item["question_content"], 
+        "metadata": {}
+    } for item in data[split]]
+    return queries
+
+def get_corpus(hf_name: str, cache_dir: str) -> list[dict]:
+    dataset = load_dataset(hf_name, cache_dir=cache_dir)["train"]
+    corpus = [
+        {"_id": i, "text": item["text"], "title": item["title"]}
+        for i,item in enumerate(dataset)
+    ]
+    return corpus
+
+
+def main():
+    dataset = datasets.load_dataset(args.dataset_name, cache_dir=args.cache_dir)
+
+    path = os.path.join(args.output_dir, args.output_name)
+    os.makedirs(path, exist_ok=True)
+    os.makedirs(os.path.join(path, "qrels"), exist_ok=True)
+
+    queries = get_queries(dataset, split="test")
+    save_file_jsonl(queries, os.path.join(path, "queries.jsonl"))
+
+    docs = get_corpus(args.corpus_name, args.cache_dir)
+    save_file_jsonl(docs, os.path.join(path, "corpus.jsonl"))
+
+    qrels = []  # no ground-truth solutions
+    qrels_path = os.path.join(path, "qrels", "test.tsv")
+    save_tsv_dict(qrels, qrels_path, ["query-id", "corpus-id", "score"])
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset_name", type=str, default="livecodebench/code_generation")
+    parser.add_argument("--corpus_name", type=str, default="code-rag-bench/programming-solutions")
+    parser.add_argument("--cache_dir", type=str, default="/scratch/zhiruow/data")
+    parser.add_argument("--output_name", type=str, default="livecodebench")
+    parser.add_argument("--output_dir", type=str, default="datasets")
+    args = parser.parse_args()
+
+    main()

research/BGE_Coder/evaluation/coderag_eval/test/create/mbpp.py

@@ -0,0 +1,78 @@
+import os
+import argparse
+import datasets
+from tqdm import tqdm
+from create.utils import save_tsv_dict, save_file_jsonl
+
+
+def get_function_name(code: str) -> str:
+    """Parse the function name for a code snippet string."""
+    lines = code.split('\n')
+    for line in lines:
+        if line.lstrip().startswith("def "):
+            break
+    func_name = line.lstrip()[4: ]
+    func_name = func_name.split('(')[0]
+    return func_name
+
+
+def document2code(data, split="test"):
+    data = data[split]
+    queries, docs, qrels = [], [], []
+
+    for item in tqdm(data):
+        doc = item["text"]
+        code = "# " + item["text"] + '\n' + item["code"]
+        doc_id = "{task_id}_doc".format_map(item)
+        code_id = "{task_id}_code".format_map(item)
+
+        queries.append({"_id": doc_id, "text": doc, "metadata": {}})
+        docs.append({"_id": code_id, "title": get_function_name(item["code"]), "text": code, "metadata": {}})
+        qrels.append({"query-id": doc_id, "corpus-id": code_id, "score": 1})
+    
+    return queries, docs, qrels
+
+
+def main():
+    dataset = datasets.load_dataset(args.dataset_name)
+
+    path = os.path.join(args.output_dir, args.output_name)
+    os.makedirs(path, exist_ok=True)
+    os.makedirs(os.path.join(path, "qrels"), exist_ok=True)
+
+    docs, queries = [], []
+    for split in args.splits:
+        queries_split, docs_split, qrels_split = document2code(dataset, split)
+        docs += docs_split
+        queries += queries_split
+
+        qrels_path = os.path.join(path, "qrels", f"{split}.tsv")
+        save_tsv_dict(qrels_split, qrels_path, ["query-id", "corpus-id", "score"])
+
+        # create canonical file for test split if not existent yet
+        if split == "test" and (not os.path.exists(args.canonical_file)):
+            canonical_solutions = []
+            for doc in docs_split:
+                canonical_solutions.append([{
+                    "text": doc["text"], "title": doc["title"]
+                }])
+            canonical_dataset = dataset["test"].add_column("docs", canonical_solutions)
+            canonical_dataset.to_json(args.canonical_file)
+    
+    save_file_jsonl(queries, os.path.join(path, "queries.jsonl"))
+    save_file_jsonl(docs, os.path.join(path, "corpus.jsonl"))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset_name", type=str, default="google-research-datasets/mbpp")
+    # parser.add_argument("--dataset_name", type=str, default="code-rag-bench/mbpp")
+    parser.add_argument("--splits", type=str, default=["train", "validation", "test"],
+                        choices=["train", "validation", "test", "prompt"])
+    parser.add_argument("--output_name", type=str, default="mbpp")
+    parser.add_argument("--output_dir", type=str, default="datasets")
+    parser.add_argument("--canonical_file", type=str, 
+                        default="datasets/canonical/mbpp_solutions.json")
+    
+    args = parser.parse_args()
+    main()

research/BGE_Coder/evaluation/coderag_eval/test/create/odex.py

@@ -0,0 +1,76 @@
+import os
+import re
+import random
+import argparse
+import datasets
+from tqdm import tqdm
+from collections import Counter
+from datasets import load_dataset
+from create.utils import save_tsv_dict, save_file_jsonl
+
+
+def document2code(data, split="test"):
+    data = data[split]
+    queries, docs, qrels = [], [], []
+
+    # build doc corpus
+    code_docs = load_dataset("neulab/docprompting-conala", "docs")["train"]
+    for i in range(len(code_docs)):
+        docs.append({
+            "_id": str(i),
+            "title": code_docs[i]["doc_id"],
+            "text": code_docs[i]["doc_content"],
+            "metadata": {}
+        })
+    
+    # load canonical docs
+    odex = load_dataset("json", data_files={"test": args.canonical_file})["test"]
+    # collect queries and query-doc matching
+    for idx,item in enumerate(tqdm(data)):
+        query = item["intent"]
+        query_id = f"{idx}_{item['task_id']}"
+        queries.append({"_id": query_id, "text": query, "metadata": {}})
+
+        doc_ids = [doc["title"] for doc in odex[idx]["docs"]]
+        for doc_id in doc_ids:
+            corpus_id = code_docs["doc_id"].index(doc_id)
+            corpus_id = str(corpus_id)
+            qrels.append({"query-id": query_id, "corpus-id": corpus_id, "score": 1})
+    
+    return queries, docs, qrels
+
+
+def main():
+    if '_' in args.dataset_name:
+        dataset_name = args.dataset_name.split('_')[0]
+        language = args.dataset_name.split('_')[1]
+    else:
+        dataset_name = args.dataset_name
+        language = 'en'
+    dataset = datasets.load_dataset(dataset_name, language) # english version by default
+
+    path = os.path.join(args.output_dir, args.output_name.replace('en', language))
+    os.makedirs(path, exist_ok=True)
+    os.makedirs(os.path.join(path, "qrels"), exist_ok=True)
+
+    docs, queries = [], []
+    for split in ["test"]:
+        queries_split, docs_split, qrels_split = document2code(dataset, split)
+        docs += docs_split
+        queries += queries_split
+
+        save_tsv_dict(qrels_split, os.path.join(path, "qrels", "{}.tsv".format(split)), ["query-id", "corpus-id", "score"])
+    
+    save_file_jsonl(queries, os.path.join(path, "queries.jsonl"))
+    save_file_jsonl(docs, os.path.join(path, "corpus.jsonl"))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset_name", type=str, default="neulab/odex")
+    parser.add_argument("--output_name", type=str, default="odex_en")
+    parser.add_argument("--canonical_file", type=str, default="datasets/canonical/odex_docs.json")
+    parser.add_argument("--output_dir", type=str, default="datasets")
+    args = parser.parse_args()
+
+    main()

research/BGE_Coder/evaluation/coderag_eval/test/create/repoeval.py

@@ -0,0 +1,306 @@
+import io
+import os
+import glob
+import json
+import argparse
+import requests
+import zipfile
+from collections import defaultdict
+from create.utils import save_tsv_dict, save_file_jsonl
+
+REPOs_line_and_api = [
+    'huggingface_diffusers',
+    'nerfstudio-project_nerfstudio',
+    'awslabs_fortuna',
+    'huggingface_evaluate',
+    'google_vizier',
+    'alibaba_FederatedScope',
+    'pytorch_rl',
+    'opendilab_ACE',
+]
+
+REPOs_function = [
+    "amazon-science_patchcore-inspection",
+    "deepmind_tracr",
+    "facebookresearch_omnivore",
+    "google_lightweight_mmm",
+    "lucidrains_imagen-pytorch",
+    "maxhumber_redframes",
+]
+
+REPO_DIRs = {
+    "api": "repositories/line_and_api_level",
+    "line": "repositories/line_and_api_level",
+    "function": "repositories/function_level",
+}
+
+
+def iterate_repository(base_dir: str, repo: str) -> dict:
+    pattern = os.path.join(f'{base_dir}/{repo}', "**", "*.py")
+    files = glob.glob(pattern, recursive=True)
+
+    skipped_files = []
+    loaded_code_files = dict()
+    base_dir_list = os.path.normpath(base_dir).split(os.sep)
+    for fname in files:
+        try:
+            code = open(fname, 'r', encoding='utf8').read()
+            fpath_tuple = tuple(os.path.normpath(fname).split(os.sep)[len(base_dir_list):])
+            loaded_code_files[fpath_tuple]= code
+        except Exception as e:
+            skipped_files.append((fname, e))
+            continue
+
+    if len(skipped_files) > 0:
+        print(f"Skipped {len(skipped_files)} out of {len(files)} files due to I/O errors")
+        for fname, e in skipped_files:
+            print(f"{fname}: {e}")
+    return loaded_code_files
+
+
+def window_overlap(span: tuple, target_span: tuple) -> bool:
+    if span[0] >= target_span[1] or span[1] <= target_span[0]:
+        return False
+    return True
+
+
+class RepoWindowMaker:
+    def __init__(self, base_dir, repo, tasks, window_size, slice_size):
+        self.base_dir = base_dir
+        self.repo = repo
+        self.window_size = window_size
+        self.slice_size = slice_size
+        self.slice_step = 1 if window_size // slice_size == 0 else window_size // slice_size
+        self.tasks = tasks
+        self.source_code_files = iterate_repository(base_dir, repo)
+        
+    def _buid_windows_for_a_file(self, fpath_tuple, code):
+        code_windows = []
+        code_lines = code.splitlines()
+        delta_size = self.window_size // 2
+        for line_no in range(0, len(code_lines), self.slice_step): # line_no starts from 0
+            start_line_no = max(0, line_no - delta_size)
+            end_line_no = min(len(code_lines), line_no + self.window_size - delta_size)
+            window_lines = [i for i in code_lines[start_line_no:end_line_no]]
+            if not window_lines:  # all empty lines
+                continue
+            window_text = '\n'.join(window_lines)
+            code_windows.append({
+                'context': window_text,
+                'metadata': {
+                    'fpath_tuple': fpath_tuple,
+                    'line_no': line_no,
+                    'start_line_no': start_line_no,
+                    'end_line_no': end_line_no,
+                    'window_size': self.window_size,
+                    'repo': self.repo,
+                    'slice_size': self.slice_size,
+                }
+            })
+        return code_windows
+    
+    def _merge_windows_with_same_context(self, code_windows):
+        merged_code_windows = defaultdict(list)
+        for code_window in code_windows:
+            context = code_window['context']
+            metadata = code_window['metadata']
+            merged_code_windows[context].append(metadata)
+        json_lines = []
+        for context, metadata_list in merged_code_windows.items():
+            json_lines.append({
+                'context': context,
+                'metadata': metadata_list
+            })
+        return json_lines
+
+    def build_windows(self):
+        all_code_windows = []
+        for fpath_tuple, code in self.source_code_files.items():
+            all_code_windows += self._buid_windows_for_a_file(fpath_tuple, code)
+        merged_code_windows = self._merge_windows_with_same_context(all_code_windows)
+        print(f'build {len(merged_code_windows)} windows for {self.repo} with window size {self.window_size} and slice {self.slice_size}')
+        ground_truth_indices = {}
+        for task in self.tasks:
+            fpath_tuple = tuple(task['metadata']['fpath_tuple'])
+            line_no = task['metadata']['line_no']
+            start_line_no = task['metadata']['context_start_lineno']
+            for i, window in enumerate(merged_code_windows):
+                if window["metadata"][0]["fpath_tuple"] != fpath_tuple:
+                    continue
+                if any([
+                    window_overlap(
+                        (sub_window["start_line_no"], sub_window["end_line_no"]), 
+                        (start_line_no, line_no + 1)
+                    )
+                    for sub_window in window["metadata"]
+                ]):
+                    if i not in ground_truth_indices: 
+                        ground_truth_indices[i] = []
+                    ground_truth_indices[i].append(task["metadata"]["task_id"])
+                    
+        return merged_code_windows, ground_truth_indices
+
+
+def download_data(directory: str = "repoeval"):
+    os.makedirs(directory, exist_ok=True)
+    
+    datasets_dir = os.path.join(directory, "datasets")
+    repos_lineapi_dir = os.path.join(directory, "repositories", "line_and_api_level")
+    repos_function_dir = os.path.join(directory, "repositories", "function_level")
+
+    print(f"Start downloading the necessary `datasets` and `repositories` files.")
+    if not os.path.exists(datasets_dir):
+        print(f"Start downloading the `datasets`.")
+        datasets_url = "https://github.com/microsoft/CodeT/raw/main/RepoCoder/datasets/datasets.zip"
+        r = requests.get(datasets_url, stream=True)
+        z = zipfile.ZipFile(io.BytesIO(r.content))
+        z.extractall(datasets_dir)
+        print("Finished downloading the `datasets` files.")
+
+    if not os.path.exists(repos_lineapi_dir):
+        print(f"Start downloading the `repositories` (line_and_api).")
+        repos_lineapi_url = "https://github.com/microsoft/CodeT/raw/main/RepoCoder/repositories/line_and_api_level.zip"
+        r = requests.get(repos_lineapi_url, stream=True)
+        z = zipfile.ZipFile(io.BytesIO(r.content))
+        z.extractall(repos_lineapi_dir)
+    
+    if not os.path.exists(repos_function_dir):
+        print(f"Start downloading the `repositories` (function).")
+        # repos_function_url = "https://github.com/microsoft/CodeT/raw/main/RepoCoder/repositories/function_level.zip"
+        repos_function_url = "https://github.com/Veronicium/repoeval_debug/raw/main/function_level.zip"
+        r = requests.get(repos_function_url, stream=True)
+        z = zipfile.ZipFile(io.BytesIO(r.content))
+        z.extractall(repos_function_dir)
+        print("Finished downloading the `repositories` files.")
+
+
+def repo2code(
+    repo: str, data_cache_dir: str, 
+    split: str, context_length: str,
+    window_size: int, slice_size: int
+):
+    # load test examples
+    file_name = f"{split}_level_completion_{context_length}_context_codex.test.jsonl"
+    if split == 'function':
+        file_name = file_name.replace('.test.jsonl', '.test.clean.jsonl')
+    
+    task_path = os.path.join(data_cache_dir, "datasets", file_name)
+    tasks = [json.loads(l.rstrip()) for l in open(task_path, 'r')]
+    tasks = [task for task in tasks if repo == task['metadata']['task_id'].split('/')[0]]
+
+    # collect queries
+    queries = []
+    for task in tasks:
+        query_id = task["metadata"]["task_id"]
+        # text = '\n'.join(task["prompt"].split('\n')[-2:])
+        text = task["prompt"]
+        metadata = task["metadata"]
+        queries.append({"_id": query_id, "text": text, "metadata": metadata})
+    
+    base_dir = os.path.join(data_cache_dir, REPO_DIRs[split])
+    repo_window_maker = RepoWindowMaker(base_dir, repo, tasks, window_size, slice_size)
+    windows, ground_truth_indices = repo_window_maker.build_windows()
+    corpus, qrels = [], []
+    query_id2gt = {task['metadata']['task_id']:[] for task in tasks}
+    for i, window in enumerate(windows):
+        path = '-'.join(window["metadata"][0]["fpath_tuple"])
+        line = f"{window['metadata'][0]['start_line_no']}-{window['metadata'][-1]['end_line_no']}"
+        corpus_id = f"{repo}_{path}_{line}"
+        corpus.append({
+            "_id": corpus_id, "title": path, 
+            "text": window["context"], "metadata": window["metadata"]
+        })
+        if i in ground_truth_indices:
+            for query_id in ground_truth_indices[i]:
+                qrels.append({"query-id": query_id, "corpus-id": corpus_id, "score": 1})
+                query_id2gt[query_id].append({"title": corpus_id.replace('_', '/'), "text": window["context"]})
+
+    return queries, corpus, qrels, query_id2gt
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--output_dir", type=str, default="datasets")
+    parser.add_argument("--results_dir", type=str, default="results")
+    parser.add_argument("--split", type=str, required=True, choices=["api", "line", "function"])
+    parser.add_argument("--context_length", type=str, default="1k", choices=["1k", "2k", "4k"])
+    parser.add_argument("--data_cache_dir", type=str, default="output/repoeval")
+    parser.add_argument("--window_size", type=int, default=20)
+    parser.add_argument("--slice_size", type=int, default=2)
+    args = parser.parse_args()
+
+    download_data(args.data_cache_dir)
+
+    path = os.path.join(args.output_dir, "repoeval", args.split)
+    os.makedirs(path, exist_ok=True)
+    REPOs = REPOs_function if args.split == "function" else REPOs_line_and_api
+    
+    file_name = f"{args.split}_level_completion_{args.context_length}_context_codex.test.jsonl"
+    data_path = os.path.join(args.data_cache_dir, "datasets", file_name)
+    data = [json.loads(l.rstrip()) for l in open(data_path, 'r')]
+    
+    # preprocess function completion data (the data in the RepoCoder repo isn't correctly formatted)
+    if args.split == 'function':        
+        repo2idx = {}
+        for task in data:
+            repo = task['metadata']['task_id'].replace('--', '_').split('/')[0]
+            if repo not in repo2idx:
+                repo2idx[repo] = 0
+            task['metadata']['task_id'] = task['metadata']['task_id'].replace('--', '_').replace('idx', str(repo2idx[repo]))
+            task['metadata']['line_no'] = task['metadata']['lineno']
+            repo2idx[repo] += 1
+            
+        new_data_path = data_path.replace('.test.jsonl', '.test.clean.jsonl')
+        with open(new_data_path, 'w') as f:
+            for task in data:
+                repo = task['metadata']['task_id'].split('/')[0]
+                if repo not in REPOs:
+                    continue
+                f.write(json.dumps(task) + '\n')
+                
+        data = [json.loads(l.rstrip()) for l in open(new_data_path, 'r')]
+
+    # build query, docs, and qrels for each repository
+    queries, corpus, qrels = [], [], []
+    query_id2gt = {}
+    for repo in REPOs:
+        repo_queries, repo_corpus, repo_qrels, repo_query_id2gt = repo2code(
+            repo, args.data_cache_dir, 
+            args.split, args.context_length,
+            args.window_size, args.slice_size
+        )
+        queries += repo_queries
+        corpus += repo_corpus
+        qrels += repo_qrels
+        query_id2gt.update(repo_query_id2gt)
+
+    save_file_jsonl(queries, os.path.join(path, "queries.jsonl"))
+    save_file_jsonl(corpus, os.path.join(path, "corpus.jsonl"))
+    save_tsv_dict(qrels, os.path.join(path, "qrels", "test.tsv"), ["query-id", "corpus-id", "score"])
+    
+    gt_data = []
+    for example in data:
+        query_id = example['metadata']['task_id']
+        gt = query_id2gt[query_id]
+        new_example = {
+            "prompt": example["prompt"],
+            "reference": example["metadata"]["ground_truth"],
+            "docs": gt[:10],
+            "metadata": {k:v for k,v in example["metadata"].items() if k != "ground_truth"},
+        }
+        gt_data.append(new_example)
+        
+    results_file = os.path.join(args.results_dir, f"repoeval-{args.split}-{args.context_length}-gt.jsonl")
+    with open(results_file, "w") as fw:
+        for ex in gt_data:
+            fw.write(json.dumps(ex) + "\n")
+        
+    results_file = os.path.join(args.results_dir, f"repoeval-{args.split}-{args.context_length}-infile.jsonl")
+    with open(results_file, "w") as fw:
+        for ex in gt_data:
+            ex = {k:v for k,v in ex.items() if k != "docs"}
+            ex["docs"] = []
+            fw.write(json.dumps(ex) + "\n")
+
+if __name__ == "__main__":
+    main()

research/BGE_Coder/evaluation/coderag_eval/test/create/repoeval_repo.py

@@ -0,0 +1,306 @@
+import io
+import os
+import glob
+import json
+import argparse
+import requests
+import zipfile
+from collections import defaultdict
+from create.utils import save_tsv_dict, save_file_jsonl
+
+REPOs_line_and_api = [
+    'huggingface_diffusers',
+    'nerfstudio-project_nerfstudio',
+    'awslabs_fortuna',
+    'huggingface_evaluate',
+    'google_vizier',
+    'alibaba_FederatedScope',
+    'pytorch_rl',
+    'opendilab_ACE',
+]
+
+REPOs_function = [
+    "amazon-science_patchcore-inspection",
+    "deepmind_tracr",
+    "facebookresearch_omnivore",
+    "google_lightweight_mmm",
+    "lucidrains_imagen-pytorch",
+    "maxhumber_redframes",
+]
+
+REPO_DIRs = {
+    "api": "repositories/line_and_api_level",
+    "line": "repositories/line_and_api_level",
+    "function": "repositories/function_level",
+}
+
+
+def iterate_repository(base_dir: str, repo: str) -> dict:
+    pattern = os.path.join(f'{base_dir}/{repo}', "**", "*.py")
+    files = glob.glob(pattern, recursive=True)
+
+    skipped_files = []
+    loaded_code_files = dict()
+    base_dir_list = os.path.normpath(base_dir).split(os.sep)
+    for fname in files:
+        try:
+            code = open(fname, 'r', encoding='utf8').read()
+            fpath_tuple = tuple(os.path.normpath(fname).split(os.sep)[len(base_dir_list):])
+            loaded_code_files[fpath_tuple]= code
+        except Exception as e:
+            skipped_files.append((fname, e))
+            continue
+
+    if len(skipped_files) > 0:
+        print(f"Skipped {len(skipped_files)} out of {len(files)} files due to I/O errors")
+        for fname, e in skipped_files:
+            print(f"{fname}: {e}")
+    return loaded_code_files
+
+
+def window_overlap(span: tuple, target_span: tuple) -> bool:
+    if span[0] >= target_span[1] or span[1] <= target_span[0]:
+        return False
+    return True
+
+
+class RepoWindowMaker:
+    def __init__(self, base_dir, repo, tasks, window_size, slice_size):
+        self.base_dir = base_dir
+        self.repo = repo
+        self.window_size = window_size
+        self.slice_size = slice_size
+        self.slice_step = 1 if window_size // slice_size == 0 else window_size // slice_size
+        self.tasks = tasks
+        self.source_code_files = iterate_repository(base_dir, repo)
+        
+    def _buid_windows_for_a_file(self, fpath_tuple, code):
+        code_windows = []
+        code_lines = code.splitlines()
+        delta_size = self.window_size // 2
+        for line_no in range(0, len(code_lines), self.slice_step): # line_no starts from 0
+            start_line_no = max(0, line_no - delta_size)
+            end_line_no = min(len(code_lines), line_no + self.window_size - delta_size)
+            window_lines = [i for i in code_lines[start_line_no:end_line_no]]
+            if not window_lines:  # all empty lines
+                continue
+            window_text = '\n'.join(window_lines)
+            code_windows.append({
+                'context': window_text,
+                'metadata': {
+                    'fpath_tuple': fpath_tuple,
+                    'line_no': line_no,
+                    'start_line_no': start_line_no,
+                    'end_line_no': end_line_no,
+                    'window_size': self.window_size,
+                    'repo': self.repo,
+                    'slice_size': self.slice_size,
+                }
+            })
+        return code_windows
+    
+    def _merge_windows_with_same_context(self, code_windows):
+        merged_code_windows = defaultdict(list)
+        for code_window in code_windows:
+            context = code_window['context']
+            metadata = code_window['metadata']
+            merged_code_windows[context].append(metadata)
+        json_lines = []
+        for context, metadata_list in merged_code_windows.items():
+            json_lines.append({
+                'context': context,
+                'metadata': metadata_list
+            })
+        return json_lines
+
+    def build_windows(self):
+        all_code_windows = []
+        for fpath_tuple, code in self.source_code_files.items():
+            all_code_windows += self._buid_windows_for_a_file(fpath_tuple, code)
+        merged_code_windows = self._merge_windows_with_same_context(all_code_windows)
+        print(f'build {len(merged_code_windows)} windows for {self.repo} with window size {self.window_size} and slice {self.slice_size}')
+        ground_truth_indices = {}
+        for task in self.tasks:
+            fpath_tuple = tuple(task['metadata']['fpath_tuple'])
+            line_no = task['metadata']['line_no']
+            start_line_no = task['metadata']['context_start_lineno']
+            for i, window in enumerate(merged_code_windows):
+                # print(window["metadata"][0]["fpath_tuple"], fpath_tuple)
+                if window["metadata"][0]["fpath_tuple"] != fpath_tuple and ' '.join(list(window["metadata"][0]["fpath_tuple"])) != ' '.join(list(fpath_tuple)):
+                    continue
+                # print(1)
+                if any([
+                    window_overlap(
+                        (sub_window["start_line_no"], sub_window["end_line_no"]), 
+                        (start_line_no, line_no + 1)
+                    )
+                    for sub_window in window["metadata"]
+                ]):
+                    # print('test')
+                    if i not in ground_truth_indices: 
+                        ground_truth_indices[i] = []
+                    ground_truth_indices[i].append(task["metadata"]["task_id"])
+        # sys.exit()
+        return merged_code_windows, ground_truth_indices
+
+
+def download_data(directory: str = "repoeval"):
+    os.makedirs(directory, exist_ok=True)
+    
+    datasets_dir = os.path.join(directory, "datasets")
+    repos_function_dir = os.path.join(directory, "repositories", "function_level")
+
+    print(f"Start downloading the necessary `datasets` and `repositories` files.")
+    if not os.path.exists(datasets_dir):
+        print(f"Start downloading the `datasets`.")
+        datasets_url = "https://github.com/microsoft/CodeT/raw/main/RepoCoder/datasets/datasets.zip"
+        r = requests.get(datasets_url, stream=True)
+        z = zipfile.ZipFile(io.BytesIO(r.content))
+        z.extractall(datasets_dir)
+        print("Finished downloading the `datasets` files.")
+
+    import shutil
+    shutil.rmtree(repos_function_dir)
+    if not os.path.exists(repos_function_dir):
+        print(f"Start downloading the `repositories` (function).")
+        repos_function_url = "https://github.com/microsoft/CodeT/raw/main/RepoCoder/repositories/function_level.zip"
+        # repos_function_url = "https://github.com/Veronicium/repoeval_debug/raw/main/function_level.zip"
+        r = requests.get(repos_function_url, stream=True)
+        z = zipfile.ZipFile(io.BytesIO(r.content))
+        z.extractall(repos_function_dir)
+        print("Finished downloading the `repositories` files.")
+
+
+def repo2code(
+    repo: str, tasks: list[dict], data_cache_dir: str, 
+    split: str, context_length: str,
+    window_size: int, slice_size: int
+):
+    # collect queries
+    queries = []
+    for task in tasks:
+        query_id = task["metadata"]["task_id"]
+        # text = '\n'.join(task["prompt"].split('\n')[-2:])
+        text = task["prompt"]
+        metadata = task["metadata"]
+        queries.append({"_id": query_id, "text": text, "metadata": metadata})
+    
+    base_dir = os.path.join(data_cache_dir, REPO_DIRs[split])
+    repo_window_maker = RepoWindowMaker(base_dir, repo, tasks, window_size, slice_size)
+    windows, ground_truth_indices = repo_window_maker.build_windows()
+    corpus, qrels = [], []
+    query_id2gt = {task['metadata']['task_id']:[] for task in tasks}
+    for i, window in enumerate(windows):
+        path = '-'.join(window["metadata"][0]["fpath_tuple"])
+        line = f"{window['metadata'][0]['start_line_no']}-{window['metadata'][-1]['end_line_no']}"
+        corpus_id = f"{repo}_{path}_{line}"
+        corpus.append({
+            "_id": corpus_id, "title": path, 
+            "text": window["context"], "metadata": window["metadata"]
+        })
+        # print(windows, ground_truth_indices)
+        if i in ground_truth_indices:
+            for query_id in ground_truth_indices[i]:
+                qrels.append({"query-id": query_id, "corpus-id": corpus_id, "score": 1})
+                query_id2gt[query_id].append({"title": corpus_id.replace('_', '/'), "text": window["context"]})
+
+    return queries, corpus, qrels, query_id2gt
+
+
+def main():
+    download_data(args.data_cache_dir)
+
+    REPOs = REPOs_function if args.split == "function" else REPOs_line_and_api
+    
+    file_name = f"{args.split}_level_completion_{args.context_length}_context_codex.test.jsonl"
+    data_path = os.path.join(args.data_cache_dir, "datasets", file_name)
+    data = [json.loads(l.rstrip()) for l in open(data_path, 'r')]
+
+    # preprocess function completion data (the data in the RepoCoder repo isn't correctly formatted)
+    if args.split == 'function':        
+        repo2idx = {}
+        for task in data:
+            repo = task['metadata']['task_id'].replace('--', '_').split('/')[0]
+            if repo not in repo2idx:
+                repo2idx[repo] = 0
+            task['metadata']['task_id'] = task['metadata']['task_id'].replace('--', '_').replace('idx', str(repo2idx[repo]))
+            task['metadata']['line_no'] = task['metadata']['lineno']
+            repo2idx[repo] += 1
+            
+        new_data_path = data_path.replace('.test.jsonl', '.test.clean.jsonl')
+        with open(new_data_path, 'w') as f:
+            for task in data:
+                repo = task['metadata']['task_id'].split('/')[0]
+                if repo not in REPOs:
+                    continue
+                f.write(json.dumps(task) + '\n')
+                
+        data = [json.loads(l.rstrip()) for l in open(new_data_path, 'r')]
+
+    # group data instances by repository
+    data_dict = {}
+    for ex in data:
+        repo_name = ex["metadata"]["task_id"]
+        repo_name = repo_name.split('/')[0]
+        if repo_name not in data_dict:
+            data_dict[repo_name] = []
+        data_dict[repo_name].append(ex)
+
+    # build query, docs, and qrels for each repository
+    for repo in REPOs:
+        queries, corpus, qrels, query_id2gt = repo2code(
+            repo, data_dict[repo], args.data_cache_dir, 
+            args.split, args.context_length,
+            args.window_size, args.slice_size
+        )
+        print(len(queries))
+        if len(qrels) == 0:
+            print(repo)
+            # sys.exit()
+            continue
+        # sys.exit()
+        path = os.path.join(args.output_dir, f"repoeval__{repo}")
+        os.makedirs(path, exist_ok=True)
+        save_file_jsonl(queries, os.path.join(path, "queries.jsonl"))
+        save_file_jsonl(corpus, os.path.join(path, "corpus.jsonl"))
+        save_tsv_dict(qrels, os.path.join(path, "qrels", "test.tsv"), ["query-id", "corpus-id", "score"])
+        
+        gt_data = []
+        for example in data_dict[repo]:
+            query_id = example['metadata']['task_id']
+            gt = query_id2gt[query_id]
+            new_example = {
+                "prompt": example["prompt"],
+                "reference": example["metadata"]["ground_truth"],
+                "docs": gt[:10],
+                "metadata": {k:v for k,v in example["metadata"].items() if k != "ground_truth"},
+            }
+            gt_data.append(new_example)
+        
+        os.makedirs(args.results_dir, exist_ok=True)
+            
+        results_file = os.path.join(args.results_dir, f"repoeval-{args.split}-{repo}-{args.context_length}-gt.jsonl")
+        with open(results_file, "w") as fw:
+            for ex in gt_data:
+                fw.write(json.dumps(ex) + "\n")
+            
+        results_file = os.path.join(args.results_dir, f"repoeval-{args.split}-{repo}-{args.context_length}-infile.jsonl")
+        with open(results_file, "w") as fw:
+            for ex in gt_data:
+                ex = {k:v for k,v in ex.items() if k != "docs"}
+                ex["docs"] = []
+                fw.write(json.dumps(ex) + "\n")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--output_dir", type=str, default="datasets")
+    parser.add_argument("--results_dir", type=str, default="results")
+    parser.add_argument("--split", type=str, default="function", choices=["function"])
+    parser.add_argument("--context_length", type=str, default="2k", choices=["1k", "2k", "4k"])
+    parser.add_argument("--data_cache_dir", type=str, default="output/repoeval")
+    parser.add_argument("--window_size", type=int, default=50)
+    parser.add_argument("--slice_size", type=int, default=5)
+    args = parser.parse_args()
+
+    main()

research/BGE_Coder/evaluation/coderag_eval/test/create/swebench.py

@@ -0,0 +1,247 @@
+import os
+import re
+import chardet
+import unidiff
+import argparse
+import datasets
+import traceback
+import subprocess
+from git import Repo
+from tqdm import tqdm
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from create.utils import save_tsv_dict, save_file_jsonl
+
+# %% Get oracle file contents
+
+# get oracle file contents from the repo
+class ContextManager:
+    def __init__(self, repo_path, base_commit, verbose=False):
+        self.repo_path = Path(repo_path).resolve().as_posix()
+        self.old_dir = os.getcwd()
+        self.base_commit = base_commit
+        self.verbose = verbose
+
+    def __enter__(self):
+        os.chdir(self.repo_path)
+        cmd = f"git reset --hard {self.base_commit} && git clean -fdxq"
+        if self.verbose:
+            subprocess.run(cmd, shell=True, check=True)
+        else:
+            subprocess.run(
+                cmd,
+                shell=True,
+                check=True,
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+            )
+        return self
+
+    def get_environment(self):
+        raise NotImplementedError()  # TODO: activate conda environment and return the environment file
+
+    def get_readme_files(self):
+        files = os.listdir(self.repo_path)
+        files = list(filter(lambda x: os.path.isfile(x), files))
+        files = list(filter(lambda x: x.lower().startswith("readme"), files))
+        return files
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        os.chdir(self.old_dir)
+
+
+class AutoContextManager(ContextManager):
+    """Automatically clones the repo if it doesn't exist"""
+
+    def __init__(self, instance, root_dir=None, verbose=False, token=None):
+        if token is None:
+            token = os.environ.get("GITHUB_TOKEN", "git")
+        self.tempdir = None
+        if root_dir is None:
+            self.tempdir = TemporaryDirectory()
+            root_dir = self.tempdir.name
+        self.root_dir = root_dir
+        repo_dir = os.path.join(self.root_dir, instance["repo"].replace("/", "__"))
+        if not os.path.exists(repo_dir):
+            repo_url = (
+                f"https://{token}@github.com/swe-bench/"
+                + instance["repo"].replace("/", "__")
+                + ".git"
+            )
+            if verbose:
+                print(f"Cloning {instance['repo']} to {root_dir}")
+            Repo.clone_from(repo_url, repo_dir)
+        super().__init__(repo_dir, instance["base_commit"], verbose=verbose)
+        self.instance = instance
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.tempdir is not None:
+            self.tempdir.cleanup()
+        return super().__exit__(exc_type, exc_val, exc_tb)
+
+
+def ingest_files(filenames):
+    files_dict = dict()
+    for filename in filenames:
+        with open(filename) as f:
+            content = f.read()
+        files_dict[filename] = content
+    return files_dict
+
+def get_oracle_filenames(instance):
+    """
+    Returns the filenames that are changed in the patch
+    """
+    source_files = {
+        patch_file.source_file.split("a/", 1)[-1]
+        for patch_file in unidiff.PatchSet(instance["patch"])
+    }
+    gold_docs = set()
+    for source_file in source_files:
+        gold_docs.add(source_file)
+    return gold_docs
+
+
+# get all file contents from the repo
+def is_test(name, test_phrases=None):
+    if test_phrases is None:
+        test_phrases = ["test", "tests", "testing"]
+    words = set(re.split(r" |_|\/|\.", name.lower()))
+    return any(word in words for word in test_phrases)
+
+def list_files(root_dir, include_tests=False):
+    files = []
+    for filename in Path(root_dir).rglob("*.py"):
+        if not include_tests and is_test(filename.as_posix()):
+            continue
+        files.append(filename.relative_to(root_dir).as_posix())
+    return files
+
+def detect_encoding(filename):
+    """
+    Detect the encoding of a file
+    """
+    with open(filename, "rb") as file:
+        rawdata = file.read()
+    return chardet.detect(rawdata)["encoding"]
+
+def ingest_directory_contents(root_dir, include_tests=False):
+    files_content = {}
+    for relative_path in list_files(root_dir, include_tests=include_tests):
+        filename = os.path.join(root_dir, relative_path)
+        encoding = detect_encoding(filename)
+        if encoding is None:
+            content = "[BINARY DATA FILE]"
+        else:
+            try:
+                with open(filename, encoding=encoding) as file:
+                    content = file.read()
+            except (UnicodeDecodeError, LookupError):
+                content = "[BINARY DATA FILE]"
+        files_content[relative_path] = content
+    return files_content
+
+def get_file_contents(input_instances, verbose: bool = False, tmp_dir: str = "/scratch"):
+    orig_dir = os.getcwd()
+    with TemporaryDirectory(dir=tmp_dir if os.path.exists(tmp_dir) else "/tmp") as root_dir:
+        for instance_id, instance in tqdm(
+            input_instances.items(),
+            total=len(input_instances),
+            desc="Getting file contents",
+        ):
+            try:
+                with AutoContextManager(instance, root_dir, verbose=verbose) as cm:
+                    readmes = cm.get_readme_files()
+                    instance["readmes"] = ingest_files(readmes)
+                    instance["oracle_file_contents"] = ingest_files(get_oracle_filenames(instance))
+                    instance["file_contents"] = ingest_directory_contents(cm.repo_path)
+                    assert all([
+                        okey in instance["file_contents"] 
+                        for okey in instance["oracle_file_contents"].keys()
+                    ])
+            except Exception as e:
+                print(f"Failed on instance {instance_id}", e)
+                traceback.print_exc()
+            finally:
+                # if AutoContextManager fails to exit properly future exits will return the wrong directory
+                os.chdir(orig_dir)
+    os.chdir(orig_dir)
+
+
+# %% Get queries, docs, and qrels
+
+def document2code(data, split: str = "test"):
+    subset = data[split]
+    if args.num_examples is not None:
+        import random
+        indices = random.sample([i for i in range(len(subset))], args.num_examples)
+        subset = subset.select(indices)
+    print(subset)
+
+    # get queries for each example
+    queries = [
+        {
+            "_id": item["instance_id"],
+            "text": item["problem_statement"], 
+            "metadata": {}
+        }
+        for item in subset
+    ]
+
+    subset_dict = {x["instance_id"]: x for x in subset}
+    get_file_contents(subset_dict, tmp_dir=args.tmp_dir)
+
+    # collect all docs, i.e., code chunks from the repo
+    docs = []
+    for instance_id, instance in subset_dict.items():
+        print(f"Instance #{instance_id}: {len(instance['oracle_file_contents'])} oracle / {len(instance['file_contents'])} files")
+        for filename, content in instance["file_contents"].items():
+            docs.append({
+                "_id": f"{instance_id}_{filename}",
+                "title": filename,
+                "text": content,
+                "metadata": {},
+            })
+
+    # find ground-truth docs for each example
+    qrels = []
+    for instance_id, instance in subset_dict.items():
+        for filename, content in instance["oracle_file_contents"].items():
+            qrels.append({
+                "query-id": instance_id,
+                "corpus-id": f"{instance_id}_{filename}",
+                "score": 1
+            })
+    
+    return queries, docs, qrels
+
+
+def main():
+    dataset = datasets.load_dataset(args.dataset_name, cache_dir=args.cache_dir)
+
+    name = "swe-bench"
+    if "lite" in args.dataset_name.lower():
+        name += "-lite"
+        
+    path = os.path.join(args.output_dir, name)
+    os.makedirs(path, exist_ok=True)
+    os.makedirs(os.path.join(path, "qrels"), exist_ok=True)
+
+    queries, docs, qrels = document2code(dataset, split="test")
+    save_file_jsonl(queries, os.path.join(path, "queries.jsonl"))
+    save_file_jsonl(docs, os.path.join(path, "corpus.jsonl"))
+    qrels_path = os.path.join(path, "qrels", "test.tsv")
+    save_tsv_dict(qrels, qrels_path, ["query-id", "corpus-id", "score"])
+    
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset_name", type=str, default="princeton-nlp/SWE-bench_Lite",
+                        choices=["princeton-nlp/SWE-bench", "princeton-nlp/SWE-bench_Lite"])
+    parser.add_argument("--cache_dir", type=str, default="/scratch/zhiruow/data")
+    parser.add_argument("--tmp_dir", type=str, default="/scratch/zhiruow/tmp")
+    parser.add_argument("--output_dir", type=str, default="datasets")
+    parser.add_argument("--num_examples", type=int, default=None)
+    args = parser.parse_args()
+
+    main()

research/BGE_Coder/evaluation/coderag_eval/test/create/swebench_repo.py

@@ -0,0 +1,263 @@
+import os
+import re
+import chardet
+import unidiff
+import argparse
+import datasets
+import traceback
+import subprocess
+from git import Repo
+from tqdm import tqdm
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from create.utils import save_tsv_dict, save_file_jsonl
+
+# %% Get oracle file contents
+
+# get oracle file contents from the repo
+class ContextManager:
+    def __init__(self, repo_path, base_commit, verbose=False):
+        self.repo_path = Path(repo_path).resolve().as_posix()
+        self.old_dir = os.getcwd()
+        self.base_commit = base_commit
+        self.verbose = verbose
+
+    def __enter__(self):
+        os.chdir(self.repo_path)
+        cmd = f"git reset --hard {self.base_commit} && git clean -fdxq"
+        if self.verbose:
+            subprocess.run(cmd, shell=True, check=True)
+        else:
+            subprocess.run(
+                cmd,
+                shell=True,
+                check=True,
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+            )
+        return self
+
+    def get_environment(self):
+        raise NotImplementedError()  # TODO: activate conda environment and return the environment file
+
+    def get_readme_files(self):
+        files = os.listdir(self.repo_path)
+        files = list(filter(lambda x: os.path.isfile(x), files))
+        files = list(filter(lambda x: x.lower().startswith("readme"), files))
+        return files
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        os.chdir(self.old_dir)
+
+
+class AutoContextManager(ContextManager):
+    """Automatically clones the repo if it doesn't exist"""
+
+    def __init__(self, instance, root_dir=None, verbose=False, token=None):
+        if token is None:
+            token = os.environ.get("GITHUB_TOKEN", "git")
+        self.tempdir = None
+        if root_dir is None:
+            self.tempdir = TemporaryDirectory()
+            root_dir = self.tempdir.name
+        self.root_dir = root_dir
+        repo_dir = os.path.join(self.root_dir, instance["repo"].replace("/", "__"))
+        if not os.path.exists(repo_dir):
+            repo_url = (
+                f"https://{token}@github.com/swe-bench/"
+                + instance["repo"].replace("/", "__")
+                + ".git"
+            )
+            if verbose:
+                print(f"Cloning {instance['repo']} to {root_dir}")
+            Repo.clone_from(repo_url, repo_dir)
+        super().__init__(repo_dir, instance["base_commit"], verbose=verbose)
+        self.instance = instance
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.tempdir is not None:
+            self.tempdir.cleanup()
+        return super().__exit__(exc_type, exc_val, exc_tb)
+
+
+def ingest_files(filenames):
+    files_dict = dict()
+    for filename in filenames:
+        with open(filename) as f:
+            content = f.read()
+        files_dict[filename] = content
+    return files_dict
+
+def get_oracle_filenames(instance):
+    """
+    Returns the filenames that are changed in the patch
+    """
+    source_files = {
+        patch_file.source_file.split("a/", 1)[-1]
+        for patch_file in unidiff.PatchSet(instance["patch"])
+    }
+    gold_docs = set()
+    for source_file in source_files:
+        gold_docs.add(source_file)
+    return gold_docs
+
+
+# get all file contents from the repo
+def is_test(name, test_phrases=None):
+    if test_phrases is None:
+        test_phrases = ["test", "tests", "testing"]
+    words = set(re.split(r" |_|\/|\.", name.lower()))
+    return any(word in words for word in test_phrases)
+
+def list_files(root_dir, include_tests=False):
+    files = []
+    for filename in Path(root_dir).rglob("*.py"):
+        if not include_tests and is_test(filename.as_posix()):
+            continue
+        files.append(filename.relative_to(root_dir).as_posix())
+    return files
+
+def detect_encoding(filename):
+    """
+    Detect the encoding of a file
+    """
+    with open(filename, "rb") as file:
+        rawdata = file.read()
+    return chardet.detect(rawdata)["encoding"]
+
+def ingest_directory_contents(root_dir, include_tests=False):
+    files_content = {}
+    for relative_path in list_files(root_dir, include_tests=include_tests):
+        filename = os.path.join(root_dir, relative_path)
+        encoding = detect_encoding(filename)
+        if encoding is None:
+            content = "[BINARY DATA FILE]"
+        else:
+            try:
+                with open(filename, encoding=encoding) as file:
+                    content = file.read()
+            except (UnicodeDecodeError, LookupError):
+                content = "[BINARY DATA FILE]"
+        files_content[relative_path] = content
+    return files_content
+
+def get_file_contents(input_instances, verbose: bool = False, tmp_dir: str = "/scratch"):
+    orig_dir = os.getcwd()
+    with TemporaryDirectory(dir=tmp_dir if os.path.exists(tmp_dir) else "/tmp") as root_dir:
+        for instance_id, instance in tqdm(
+            input_instances.items(),
+            total=len(input_instances),
+            desc="Getting file contents",
+        ):
+            try:
+                with AutoContextManager(instance, root_dir, verbose=verbose) as cm:
+                    readmes = cm.get_readme_files()
+                    instance["readmes"] = ingest_files(readmes)
+                    instance["oracle_file_contents"] = ingest_files(get_oracle_filenames(instance))
+                    instance["file_contents"] = ingest_directory_contents(cm.repo_path)
+                    assert all([
+                        okey in instance["file_contents"] 
+                        for okey in instance["oracle_file_contents"].keys()
+                    ])
+            except Exception as e:
+                print(f"Failed on instance {instance_id}", e)
+                traceback.print_exc()
+            finally:
+                # if AutoContextManager fails to exit properly future exits will return the wrong directory
+                os.chdir(orig_dir)
+    os.chdir(orig_dir)
+
+
+import multiprocessing as mp
+from functools import partial
+
+def process_single_item(item, args):
+    """处理单个数据项的函数"""
+    name = "swe-bench"
+    if "lite" in args.dataset_name.lower():
+        name += "-lite"
+        
+    queries = [{
+        "_id": item["instance_id"],
+        "text": item["problem_statement"],
+        "metadata": {}
+    }]
+    item_dict = {item["instance_id"]: item}
+    
+    output_path = os.path.join(args.output_dir, f"{name}_{item['instance_id']}", "qrels", "test.tsv")
+    if os.path.exists(output_path):
+        return
+        
+    try:
+        get_file_contents(item_dict, tmp_dir=args.tmp_dir)
+        
+        docs = []
+        for instance_id, instance in item_dict.items():
+            print(f"Instance #{instance_id}: {len(instance['oracle_file_contents'])} oracle / {len(instance['file_contents'])} files")
+            for filename, content in instance["file_contents"].items():
+                docs.append({
+                    "_id": f"{instance_id}_{filename}",
+                    "title": filename,
+                    "text": content,
+                    "metadata": {},
+                })
+
+        qrels = []
+        for instance_id, instance in item_dict.items():
+            for filename, content in instance["oracle_file_contents"].items():
+                qrels.append({
+                    "query-id": instance_id,
+                    "corpus-id": f"{instance_id}_{filename}",
+                    "score": 1
+                })
+
+        path = os.path.join(args.output_dir, f"{name}_{instance_id}")
+        os.makedirs(path, exist_ok=True)
+        os.makedirs(os.path.join(path, "qrels"), exist_ok=True)
+        
+        save_file_jsonl(queries, os.path.join(path, "queries.jsonl"))
+        save_file_jsonl(docs, os.path.join(path, "corpus.jsonl"))
+        qrels_path = os.path.join(path, "qrels", "test.tsv")
+        save_tsv_dict(qrels, qrels_path, ["query-id", "corpus-id", "score"])
+        
+    except Exception as e:
+        print(f"Error processing item {item['instance_id']}: {str(e)}")
+
+def main():
+    dataset = datasets.load_dataset(args.dataset_name, cache_dir=args.cache_dir)["test"]
+    if args.num_examples is not None:
+        import random
+        indices = random.sample([i for i in range(len(dataset))], args.num_examples)
+        dataset = dataset.select(indices)
+    print(dataset)
+
+    # 创建进程池
+    num_processes = mp.cpu_count() - 1  # 留一个CPU核心
+    pool = mp.Pool(processes=num_processes)
+    
+    # 使用partial固定args参数
+    process_func = partial(process_single_item, args=args)
+    
+    # 使用进程池并行处理
+    list(tqdm(
+        pool.imap_unordered(process_func, dataset),
+        total=len(dataset),
+        desc="Processing items"
+    ))
+    
+    # 关闭进程池
+    pool.close()
+    pool.join()
+    
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset_name", type=str, default="princeton-nlp/SWE-bench_Lite",
+                        choices=["princeton-nlp/SWE-bench", "princeton-nlp/SWE-bench_Lite"])
+    parser.add_argument("--cache_dir", type=str, default="/scratch/zhiruow/data")
+    parser.add_argument("--tmp_dir", type=str, default="/scratch/zhiruow/tmp")
+    parser.add_argument("--output_dir", type=str, default="datasets")
+    parser.add_argument("--num_examples", type=int, default=None)
+    args = parser.parse_args()
+
+    main()

research/BGE_Coder/evaluation/coderag_eval/test/create/utils.py

@@ -0,0 +1,37 @@
+import jsonlines
+import csv
+import os
+
+def load_jsonlines(file):
+    with jsonlines.open(file, 'r') as jsonl_f:
+        lst = [obj for obj in jsonl_f]
+    return lst
+
+def save_file_jsonl(data, fp):
+    with jsonlines.open(fp, mode='w') as writer:
+        writer.write_all(data)
+
+def save_tsv_dict(data, fp, fields):
+    # build dir
+    dir_path = os.path.dirname(fp)
+    os.makedirs(dir_path, exist_ok=True)
+    
+    # writing to csv file
+    with open(fp, 'w') as csvfile:
+        writer = csv.DictWriter(csvfile, fieldnames=fields, delimiter='\t',)
+        writer.writeheader()
+        writer.writerows(data)
+
+def cost_esitmate(path):
+    corpus = load_jsonlines(os.path.join(path, "corpus.jsonl"))
+    queries = load_jsonlines(os.path.join(path, "queries.jsonl"))
+    num_corpus_words = 0
+    num_queries_words = 0
+    for item in tqdm(corpus):
+        num_corpus_words += len(item["text"].split(" "))
+    for item in tqdm(queries):
+        num_queries_words += len(item["text"].split(" "))
+    print(len(corpus))
+    print(len(queries))
+    print(num_corpus_words)
+    print(num_queries_words)