chore: import upstream snapshot with attribution

2026-07-13 13:39:21 +08:00
commit bc34f6df14
1149 changed files with 328099 additions and 0 deletions
@@ -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()