chore: import upstream snapshot with attribution

research/BGE_Coder/README.md

@@ -0,0 +1,233 @@
+<h1 align="center">CodeR: Towards A Generalist Code Embedding Model</h1>
+<p align="center">
+    <a href="https://huggingface.co/datasets/nebula2025/CodeR-Pile">
+        <img alt="Build" src="https://img.shields.io/badge/🤗 Dataset-CodeR Pile-yellow">
+    </a>
+    <a href="https://huggingface.co/nebula2025/CodeR-full">
+        <img alt="Build" src="https://img.shields.io/badge/🤗 Model-CodeR Full-green">
+    </a>
+    <a href="https://huggingface.co/nebula2025/CodeR-synthetic">
+        <img alt="Build" src="https://img.shields.io/badge/🤗 Model-CodeR Synthetic-blue">
+    </a>
+</p>
+
+
+This repo contains the data, training, and evaluation pipeline for CodeR / [BGE-Code-v1](https://huggingface.co/BAAI/bge-code-v1)
+
+**[BGE-Code-v1](https://huggingface.co/BAAI/bge-code-v1)** is an LLM-based code embedding model that supports code retrieval, text retrieval, and multilingual retrieval. It primarily demonstrates the following capabilities:
+
+- Superior Code Retrieval Performance: The model demonstrates exceptional code retrieval capabilities, supporting natural language queries in both English and Chinese, as well as 20 programming languages.
+- Robust Text Retrieval Capabilities: The model maintains strong text retrieval capabilities comparable to text embedding models of similar scale.
+- Extensive Multilingual Support: BGE-Code-v1 offers comprehensive multilingual retrieval capabilities, excelling in languages such as English, Chinese, Japanese, French, and more.
+
+## :bell: News:
+
+- 🥳 5/15/2025: We have released the CodeR! :fire:
+
+## Usage
+
+### Using FlagEmbedding
+
+```
+git clone https://github.com/FlagOpen/FlagEmbedding.git
+cd FlagEmbedding
+pip install -e .
+from FlagEmbedding import FlagLLMModel
+queries = [
+    "Delete the record with ID 4 from the 'Staff' table.", 
+    'Delete all records in the "Livestock" table where age is greater than 5'
+]
+documents = [
+    "DELETE FROM Staff WHERE StaffID = 4;",
+    "DELETE FROM Livestock WHERE age > 5;"
+]
+model = FlagLLMModel('BAAI/bge-code-v1', 
+                     query_instruction_format="<instruct>{}\n<query>{}",
+                     query_instruction_for_retrieval="Given a question in text, retrieve SQL queries that are appropriate responses to the question.",
+                     trust_remote_code=True,
+                     use_fp16=True) # Setting use_fp16 to True speeds up computation with a slight performance degradation
+embeddings_1 = model.encode_queries(queries)
+embeddings_2 = model.encode_corpus(documents)
+similarity = embeddings_1 @ embeddings_2.T
+print(similarity)
+```
+
+By default, FlagLLMModel will use all available GPUs when encoding. Please set `os.environ["CUDA_VISIBLE_DEVICES"]` to select specific GPUs. You also can set `os.environ["CUDA_VISIBLE_DEVICES"]=""` to make all GPUs unavailable.
+
+### Using Sentence Transformers
+
+```python
+from sentence_transformers import SentenceTransformer
+import torch
+
+# Load the model, optionally in float16 precision for faster inference
+model = SentenceTransformer(
+    "BAAI/bge-code-v1",
+    trust_remote_code=True,
+    model_kwargs={"torch_dtype": torch.float16},
+)
+
+# Prepare a prompt given an instruction
+instruction = 'Given a question in text, retrieve SQL queries that are appropriate responses to the question.'
+prompt = f'<instruct>{instruction}\n<query>'
+# Prepare queries and documents
+queries = [
+    "Delete the record with ID 4 from the 'Staff' table.", 
+    'Delete all records in the "Livestock" table where age is greater than 5'
+]
+documents = [
+    "DELETE FROM Staff WHERE StaffID = 4;",
+    "DELETE FROM Livestock WHERE age > 5;"
+]
+
+# Compute the query and document embeddings
+query_embeddings = model.encode(queries, prompt=prompt)
+document_embeddings = model.encode(documents)
+
+# Compute the cosine similarity between the query and document embeddings
+similarities = model.similarity(query_embeddings, document_embeddings)
+print(similarities)
+```
+
+### Using HuggingFace Transformers
+
+```python
+import torch
+import torch.nn.functional as F
+
+from torch import Tensor
+from transformers import AutoTokenizer, AutoModel
+
+
+def last_token_pool(last_hidden_states: Tensor,
+                 attention_mask: Tensor) -> Tensor:
+    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
+    if left_padding:
+        return last_hidden_states[:, -1]
+    else:
+        sequence_lengths = attention_mask.sum(dim=1) - 1
+        batch_size = last_hidden_states.shape[0]
+        return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]
+
+
+def get_detailed_instruct(task_description: str, query: str) -> str:
+    return f'<instruct>{task_description}\n<query>{query}'
+
+
+instruction = 'Given a question in text, retrieve SQL queries that are appropriate responses to the question.'
+queries = [
+    "Delete the record with ID 4 from the 'Staff' table.", 
+    'Delete all records in the "Livestock" table where age is greater than 5'
+]
+documents = [
+    "DELETE FROM Staff WHERE StaffID = 4;",
+    "DELETE FROM Livestock WHERE age > 5;"
+]
+input_texts = queries + documents
+
+tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-code-v1', trust_remote_code=True)
+model = AutoModel.from_pretrained('BAAI/bge-code-v1', trust_remote_code=True)
+model.eval()
+
+max_length = 4096
+# Tokenize the input texts
+batch_dict = tokenizer(input_texts, max_length=max_length, padding=True, truncation=True, return_tensors='pt', pad_to_multiple_of=8)
+
+with torch.no_grad():
+    outputs = model(**batch_dict)
+    embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
+    
+# normalize embeddings
+embeddings = F.normalize(embeddings, p=2, dim=1)
+scores = (embeddings[:2] @ embeddings[2:].T) * 100
+print(scores.tolist())
+```
+
+## Evaluation
+
+**BGE-Code-v1** achieves state-of-the-art performance on both the CoIR and CodeRAG benchmarks.
+
+- CoIR
+
+|                       | CodeXEmbed-2B | CodeXEmbed-7B | Voyage-Code-002 | Voyage-Code-003 | BGE-Code-v1 |
+| --------------------- | ------------- | ------------- | --------------- | --------------- | ----------- |
+| **Apps**              | 76.86         | 85.38         | 26.52           | 93.62           | 98.08       |
+| **CosQA**             | 40.47         | 42.47         | 29.79           | 34.45           | 46.72       |
+| **Text2SQL**          | 78.42         | 78.94         | 69.26           | 62.87           | 64.35       |
+| **CSN**               | 87.87         | 89.67         | 81.79           | 89.35           | 89.53       |
+| **CSN-CCR**           | 97.66         | 97.95         | 73.45           | 90.05           | 98.30       |
+| **CodeTrans-Contest** | 90.30         | 94.45         | 72.77           | 94.96           | 94.38       |
+| **CodeTrans-DL**      | 38.57         | 40.46         | 27.48           | 38.57           | 46.13       |
+| **StackOverFlow-QA**  | 94.47         | 96.33         | 67.68           | 97.17           | 95.35       |
+| **CodeFeedBack-ST**   | 86.36         | 87.53         | 65.35           | 90.67           | 90.56       |
+| **CodeFeedBack-MT**   | 65.51         | 68.83         | 28.74           | 93.58           | 94.38       |
+| **AVG**               | **75.65**     | **78.20**     | **56.26**       | **78.53**       | **81.77**   |
+
+- CodedRAG
+
+|                 | HummanEval | MBPP | DS-1000 | ODEX | RepoEval | SWE-bench-Lite | AVG      |
+| --------------- | ---------- | ---- | ------- | ---- | -------- | -------------- | -------- |
+| SFR             | 100.0      | 99.0 | 19.3    | 37.1 | 83.8     | 62.7           | **67.0** |
+| Jina-v2-code    | 100.0      | 97.7 | 26.2    | 19.9 | 90.5     | 58.3           | **65.4** |
+| CodeXEmbed-2B   | 100.0      | 97.4 | 25.4    | 23.9 | 88.7     | 52.4           | **64.6** |
+| Voyage-Code-002 | 100.0      | 99.0 | 33.1    | 26.6 | 94.3     | 29.1           | **63.7** |
+| BGE-Code-v1     | 100.0      | 99.2 | 40.9    | 36.1 | 93.1     | 67.4           | **72.8** |
+
+### Instructions for Evaluation
+
+```python
+{
+    "Apps": "Given a code contest problem description, retrieve relevant code that can help solve the problem.",
+    "CosQA": "Given a web search query, retrieve relevant code that can help answer the query.",
+    "Text2SQL": "Given a question in text, retrieve SQL queries that are appropriate responses to the question.",
+    "CSN": "Given a piece of code, retrieve the document string that summarizes the code.",
+    "CSN-CCR": "Given a piece of code segment, retrieve the code segment that is the latter part of the code.",
+    "CodeTrans-DL": "Given a piece of code, retrieve code that is semantically equivalent to the input code.",
+    "CodeTrans-Contest": "Given a piece of Python code, retrieve C++ code that is semantically equivalent to the input code.",
+    "StackOverFlow-QA": "Given a question that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.",
+    "CodeFeedBack-ST": "Given a question that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.",
+    "CodeFeedBack-MT": "Given a multi-turn conversation history that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.",
+    "HummanEval": "Given a question that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.",
+    "MBPP": "Given a textual explanation of code functionality, retrieve the corresponding code implementation.",
+    "DS-1000": "Given a question that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.",
+    "ODEX": "Given a question, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.",
+    "RepoEval": "Given a piece of code segment, retrieve the code segment that is the latter part of the code.",
+    "SWE-bench-Lite": "Given a code snippet containing a bug and a natural language description of the bug or error, retrieve code snippets that demonstrate solutions or fixes for similar bugs or errors (the desired documents)."
+}
+```
+
+### Evaluation script
+
+#### CoIR
+
+For CoIR, we use the [CoIR](https://github.com/CoIR-team/coir) evaluation script:
+
+```shell
+cd ./evaluation/coir_eval
+### clone coir
+mkdir test
+cd ./test
+git clone https://github.com/CoIR-team/coir.git
+mv ./coir/coir ../
+cd ..
+rm -rf ./test
+### evaluate
+bash eval.sh
+```
+
+### CodeRAG
+
+For CodeRAG, we use the [CodeRAG](https://github.com/code-rag-bench/code-rag-bench) evaluation script:
+
+```shell
+cd ./evaluation/coderag_eval
+### clone coderag
+git clone https://github.com/code-rag-bench/code-rag-bench.git
+## You need prepare environment according to README.md
+rm -rf ./code-rag-bench/retrieval/create
+cp -r ./test/* ./code-rag-bench/retrieval/
+### prepare data
+bash prepare_data.sh
+### evaluate
+bash eval.sh
+```

research/BGE_Coder/data_generation/corpus_generator.py

@@ -0,0 +1,104 @@
+import os
+import random
+import datasets
+from tqdm import tqdm
+from typing import List, Tuple
+
+from utils import clean_code
+from constant import DocLength
+
+
+class CorpusGenerator:
+    def __init__(
+        self,
+        cache_dir: str = None,
+    ):
+        self.cache_dir = cache_dir
+
+    def _load_corpus(self, corpus_dir: str, doc_length: List[str], external_path: List[str],
+                     source_language: str, stop_threshold: int = -1):
+        """
+        Load availavle documents for a given task from the CoIR-Retrieval dataset.
+        """
+
+        corpus_list = []
+
+        if corpus_dir is not None and os.path.exists(corpus_dir):
+            file_list = os.listdir(corpus_dir)
+            random.shuffle(file_list)
+            
+            for file in file_list:
+                flag = False
+                if not file.endswith('.jsonl'):
+                    flag = False
+                for d_length in doc_length:
+                    d_length = DocLength[d_length].value
+                    if d_length in file:
+                        flag = True
+                if flag is False:
+                    continue
+                file_path = os.path.join(corpus_dir, file)
+                corpus = datasets.load_dataset('json', data_files=file_path, cache_dir=self.cache_dir)['train']
+                for data in tqdm(corpus, desc="Loading corpus"):
+                    if source_language is None:
+                        lang = os.path.basename(corpus_dir)
+                        data['language'] = lang
+                    else:
+                        data['language'] = source_language
+                    
+                    text = clean_code(data["text"], data["language"], length_threshold=200)
+                    data["text"] = text
+                    if text != '':
+                        corpus_list.append(data)
+                
+                if stop_threshold > 0 and len(corpus_list) > stop_threshold:
+                    break
+                break
+
+        for ep in external_path:
+            if os.path.exists(ep):
+                corpus = datasets.load_dataset('json', data_files=ep, cache_dir=self.cache_dir)['train']
+                for data in tqdm(corpus, desc="Loading corpus"):
+                    if source_language is None:
+                        lang = os.path.basename(os.path.dirname(ep))
+                        data['language'] = lang
+                    else:
+                        data['language'] = source_language
+                    
+                    # useful when the text is not present in the data
+                    if "text" not in data:
+                        data["text"] = data["pos"][0]
+                    
+                    corpus_list.append(data)
+                    text = clean_code(data["text"], lang, length_threshold=200)
+                    data["text"] = text
+                    if text != '':
+                        corpus_list.append(data)
+
+        return corpus_list
+
+    def run(
+        self,
+        num_samples: int = -1,
+        max_corpus: int = -1,
+        corpus_dir: str = None,
+        doc_length: List[str] = ["len_0_500"],
+        external_path: List[str] = None,
+        source_language: str = None
+    ) -> Tuple[List[dict], List[dict]]:
+        stop_threshold = max(num_samples * 10, max_corpus * 2)
+        corpus_list = self._load_corpus(
+            corpus_dir, doc_length, external_path, source_language, stop_threshold
+        )
+
+        if num_samples > 0 and num_samples < len(corpus_list):
+            small_corpus_list = random.sample(corpus_list, num_samples)
+        else:
+            small_corpus_list = corpus_list
+        
+        if max_corpus > 0 and max_corpus < len(corpus_list):
+            corpus_list = random.sample(corpus_list, max_corpus)
+        else:
+            corpus_list = corpus_list
+
+        return small_corpus_list, corpus_list

research/BGE_Coder/data_generation/format_generated_examples.py

@@ -0,0 +1,127 @@
+import os
+import json
+from constant import Language, CodeLanguage, TaskType, CODE_TRANSLATION_RETRIEVAL_PAIRS, \
+    get_pos_as_input_by_task_type
+
+
+def format_generated_examples(
+    file_path: str,
+    save_path: str,
+    task_type: TaskType
+):
+    if os.path.exists(save_path):
+        return
+    
+    if not os.path.exists(file_path):
+        print("====================================")
+        print("Warning: file not found! Maybe need to generate it first.")
+        print(f"file_path: {file_path}")
+        return
+    
+    pos_as_input = get_pos_as_input_by_task_type(task_type)
+    
+    data_list = []
+    with open(file_path, "r", encoding="utf-8") as f:
+        for line in f.readlines():
+            data = json.loads(line)
+            
+            if pos_as_input:
+                _input = data["pos"][0]
+                _output = data["query"]
+            else:
+                _input = data["query"]
+                _output = data["pos"][0]
+
+            if 'provided' in _input:
+                continue
+            if len(_input) > 12000 or len(_output) > 12000:
+                continue
+
+            data_list.append({
+                "input": _input,
+                "output": _output
+            })
+    
+    if len(data_list) == 0:
+        print("====================================")
+        print("Warning: no data found!")
+        print(f"file_path: {file_path}")
+        return
+    
+    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+    with open(save_path, "w", encoding="utf-8") as f:
+        json.dump(data_list, f, indent=4, ensure_ascii=False)
+
+
+def main():
+    original_gen_examples_dir = "./examples"
+    
+    formatted_examples_dir = "./filtered_for_generation"
+    
+    for language in Language:
+        for task_type in TaskType:
+            if task_type == TaskType.code_translation_retrieval:
+                for code_language_pair in CODE_TRANSLATION_RETRIEVAL_PAIRS:
+                    code_language, tgt_code_language = code_language_pair
+                    
+                    file_path = os.path.join(
+                        original_gen_examples_dir,
+                        language.name, task_type.name, f"{language.name}-{code_language.name}-to-{tgt_code_language.name}-triplets.jsonl"
+                    )
+                    save_path = os.path.join(
+                        formatted_examples_dir,
+                        language.name, task_type.name, f"{code_language.name}-to-{tgt_code_language.name}_sample_examples.json"
+                    )
+                    
+                    format_generated_examples(file_path, save_path, task_type)
+                    
+                for code_language_pair in CODE_TRANSLATION_RETRIEVAL_PAIRS:
+                    tgt_code_language, code_language = code_language_pair
+                    
+                    file_path = os.path.join(
+                        original_gen_examples_dir,
+                        language.name, task_type.name, f"{language.name}-{code_language.name}-to-{tgt_code_language.name}-triplets.jsonl"
+                    )
+                    save_path = os.path.join(
+                        formatted_examples_dir,
+                        language.name, task_type.name, f"{code_language.name}-to-{tgt_code_language.name}_sample_examples.json"
+                    )
+                    
+                    format_generated_examples(file_path, save_path, task_type)
+                
+            elif task_type == TaskType.text2sql_retrieval:
+                file_path = os.path.join(
+                    original_gen_examples_dir,
+                    language.name, task_type.name, f"{language.name}-sql-triplets.jsonl"
+                )
+                save_path = os.path.join(
+                    formatted_examples_dir,
+                    language.name, task_type.name, "sql_sample_examples.json"
+                )
+                
+                format_generated_examples(file_path, save_path, task_type)
+            
+            elif task_type == TaskType.code_context_retrieval:
+                continue
+            
+            else:
+                for code_language in CodeLanguage:
+                    if code_language == CodeLanguage.null:
+                        continue
+                    
+                    file_path = os.path.join(
+                        original_gen_examples_dir,
+                        language.name, task_type.name, f"{language.name}-{code_language.name}-triplets.jsonl"
+                    )
+                    save_path = os.path.join(
+                        formatted_examples_dir,
+                        language.name, task_type.name, f"{code_language.name}_sample_examples.json"
+                    )
+                    
+                    format_generated_examples(file_path, save_path, task_type)
+    
+    print("All done!")
+
+
+if __name__ == "__main__":
+    main()

research/BGE_Coder/data_generation/llm.py

@@ -0,0 +1,134 @@
+import os
+import time
+import openai
+import random
+import tiktoken
+import threading
+from openai import OpenAI, AzureOpenAI
+from typing import Tuple
+
+
+class LLM:
+    def __init__(
+        self,
+        model: str="Qwen2-5-Coder-32B-Instruct",
+        model_type: str = "open-source",
+        port: int = 8000,
+    ):
+        if model_type == "open-source":
+            self.client = OpenAI(
+                api_key="EMPTY",
+                base_url=f"http://localhost:{port}/v1/"
+            )
+        elif model_type == "azure":
+            self.client = AzureOpenAI(
+                api_key=os.getenv("OPENAI_API_KEY"),
+                api_version=os.getenv("AZURE_API_VERSION", "2024-02-01"),
+                azure_endpoint=os.getenv("AZURE_ENDPOINT"),
+                azure_deployment=os.getenv("OPENAI_DEPLOYMENT_NAME", 'gpt-35-turbo')
+            )
+        elif model_type == "openai":
+            self.client = OpenAI(
+                api_key=os.getenv("OPENAI_API_KEY"),
+                base_url=os.getenv("OPENAI_BASE_URL", None)
+            )
+        else:
+            raise ValueError("model_type must be one of ['open-source', 'azure', 'openai']")
+        
+        self.model = model
+        self.tokenizer = tiktoken.get_encoding("o200k_base")
+    
+    def split_text(self, text: str, anchor_points: Tuple[float, float] = (0.4, 0.7)):
+        token_ids = self.tokenizer.encode(text)
+        anchor_point = random.uniform(anchor_points[0], anchor_points[1])
+        split_index = int(len(token_ids) * anchor_point)
+        return self.tokenizer.decode(token_ids[:split_index]), self.tokenizer.decode(token_ids[split_index:])
+    
+    def chat(
+        self,
+        prompt: str,
+        max_tokens: int = 8192,
+        logit_bais: dict = None,
+        n: int = 1,
+        temperature: float = 1.0,
+        top_p: float = 0.6,
+        repetition_penalty: float = 1.0,
+        remove_thinking: bool = True,
+        timeout: int = 90,
+    ):
+        endure_time = 0
+        endure_time_limit = timeout * 2
+        
+        def create_completion(results):
+            try:
+                completion = self.client.chat.completions.create(
+                    model=self.model,
+                    messages=[{"role": "user", "content": prompt}],
+                    max_tokens=max_tokens,
+                    logit_bias=logit_bais if logit_bais is not None else {},
+                    n=n,
+                    temperature=temperature,
+                    top_p=top_p,
+                    extra_body={'repetition_penalty': repetition_penalty},
+                    timeout=timeout,
+                )
+                results["content"] = [x.message.content for x in completion.choices[:n]]
+            except openai.BadRequestError as e:
+                # The response was filtered due to the prompt triggering Azure OpenAI's content management policy.
+                results["content"] = [None for _ in range(n)]
+            except openai.APIConnectionError as e:
+                results["error"] = f'APIConnectionError({e})'
+            except openai.RateLimitError as e:
+                results["error"] = f'RateLimitError({e})'
+            except Exception as e:
+                results["error"] = f"Error: {e}"
+        
+        while True:
+            results = {"content": None, "error": None}
+            completion_thread = threading.Thread(target=create_completion, args=(results,))
+            completion_thread.start()
+            
+            start_time = time.time()
+            while completion_thread.is_alive():
+                elapsed_time = time.time() - start_time
+                if elapsed_time > endure_time_limit:
+                    print("Completion timeout exceeded. Aborting...")
+                    return [None for _ in range(n)]
+                time.sleep(1)
+            
+            # If an error occurred during result processing
+            if results["error"]:
+                if endure_time >= endure_time_limit:
+                    print(f'{results["error"]} - Skip this prompt.')
+                    return [None for _ in range(n)]
+                print(f"{results['error']} - Waiting for 5 seconds...")
+                endure_time += 5
+                time.sleep(5)
+                continue
+            
+            content_list = results["content"]
+            if remove_thinking:
+                content_list = [x.split('</think>')[-1].strip('\n').strip() if x is not None else None for x in content_list]
+            return content_list
+
+
+if __name__ == "__main__":
+    llm = LLM(
+        model="gpt-4o-mini-2024-07-18",
+        model_type="openai"
+    )
+
+    prompt = "hello, who are you?"
+    response = llm.chat(prompt)[0]
+    print(response)
+
+
+if __name__ == "__main__":
+    llm = LLM(
+        model="gpt-4o-mini-2024-07-18",
+        model_type="openai"
+    )
+
+    prompt = "hello, who are you?"
+    response = llm.chat(prompt)[0]
+    print(response)

research/BGE_Coder/data_generation/run_generation.py

@@ -0,0 +1,368 @@
+import os
+import json
+import time
+import gc
+import torch
+import argparse
+import random
+from hashlib import md5
+import multiprocessing as mp
+from typing import List, Optional
+
+from constant import TaskType, Language, CodeLanguage, NUM_HARD_NEGATIVES
+from corpus_generator import CorpusGenerator
+from triplet_generator import TripletGenerator
+from search import get_top1
+
+
+def compute_md5(text: str):
+    return md5(text.encode()).hexdigest()
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--task_type',
+        type=str,
+        required=True,
+        help='The task type to generate data for',
+        choices=[t.name for t in TaskType]
+    )
+    parser.add_argument(
+        '--code_language',
+        type=str,
+        required=True,
+        help='The code language to generate questions for.',
+        choices=[c.name for c in CodeLanguage]
+    )
+    parser.add_argument(
+        '--corpus_root',
+        type=str,
+        required=True,
+        help='The root directory of the corpus data.'
+    )
+    parser.add_argument(
+        '--save_dir',
+        type=str,
+        required=True,
+        help='The path to save the generated data'
+    )
+    parser.add_argument(
+        '--examples_dir',
+        type=str,
+        default=None,
+        help='The path to the examples directory. If not None, the examples will be used for few-shot generation.'
+    )
+    parser.add_argument(
+        '--num_examples',
+        type=int,
+        default=3,
+        help='The number of examples to use for few-shot generation. Default: 3'
+    )
+    parser.add_argument(
+        '--cache_dir',
+        type=str,
+        default=None,
+        help='The cache directory'
+    )
+    parser.add_argument(
+        '--language',
+        type=str,
+        default='en',
+        help='The language to generate for. ISO 639-1 code. Default: en',
+        choices=[l.name for l in Language]
+    )
+    parser.add_argument(
+        '--tgt_code_language',
+        type=str,
+        default=None,
+        help='The target code language to generate code translations for.',
+        choices=[c.name for c in CodeLanguage]
+    )
+    parser.add_argument(
+        '--num_samples',
+        type=int,
+        default=-1,
+        help='The number of examples to use for generation. Default: -1. Use all available examples.'
+    )
+    parser.add_argument(
+        '--model',
+        type=str,
+        default='Qwen2.5-72B-Instruct',
+        help='The model to use for generation. Default: Qwen2.5-72B-Instruct'
+    )
+    parser.add_argument(
+        '--model_type',
+        type=str,
+        default='open-source',
+        help='The type of model to use for generation. Default: open-source',
+    )
+    parser.add_argument(
+        '--port',
+        type=int,
+        default=8000,
+        help='The port for vllm.'
+    )
+    parser.add_argument(
+        '--num_processes',
+        type=int,
+        default=1,
+        help='The number of processes to use for generation. Default: 1'
+    )
+    parser.add_argument(
+        '--doc_length',
+        type=str,
+        default='len_0_500',
+        help='The corpus length used to load dataset. Default: len_0_500'
+    )
+    parser.add_argument(
+        '--external_path',
+        type=str,
+        default='',
+        help='The corpus length used to load dataset. Default: len_0_500'
+    )
+    parser.add_argument(
+        '--sim_model_name',
+        type=str,
+        default=None,
+        help='The language of source corpus.'
+    )
+    parser.add_argument(
+        '--max_corpus',
+        type=int,
+        default=500000,
+        help='The max num of corpus to load.'
+    )
+    parser.add_argument(
+        '--overwrite',
+        action='store_true',
+        help='Whether to overwrite the existing data.'
+    )
+    parser.add_argument(
+        '--debug_mode',
+        action='store_true',
+        help='Whether to open debug mode.'
+    )
+    parser.add_argument(
+        '--gen_hard_neg',
+        action='store_true',
+        help='Whether to generate hard negatives.'
+    )
+    parser.add_argument(
+        '--seed',
+        type=int,
+        default=None,
+        help='Random seed for generating triplets using the same positive. Default: 42'
+    )
+    args = parser.parse_args()
+    return args
+
+
+def gen_triplets(
+    model: str,
+    model_type: str,
+    port: int,
+    positives: List[dict],
+    task_type: str,
+    language: str,
+    code_language: str,
+    tgt_code_language: str,
+    examples_pool: Optional[List[dict]] = None,
+    num_examples: int = 3,
+    tqdm_desc: str = "Generating triplets",
+    thread_count: int = 1,
+    gen_cache_dir: Optional[str] = None,
+    debug_mode: bool = False,
+    gen_hard_neg: bool = False,
+):
+    triplet_generator = TripletGenerator(model, model_type, port, cache_dir=gen_cache_dir)
+    triplets = triplet_generator.run(
+        positives=positives,
+        task_type=task_type,
+        language=language,
+        code_language=code_language,
+        tgt_code_language=tgt_code_language,
+        examples_pool=examples_pool,
+        num_examples=num_examples,
+        tqdm_desc=tqdm_desc,
+        thread_count=thread_count,
+        debug_mode=debug_mode,
+        gen_hard_neg=gen_hard_neg,
+        num_negatives=NUM_HARD_NEGATIVES,
+    )
+    return triplets
+
+
+def get_save_path(
+    save_dir: str,
+    task_type: str,
+    language: str,
+    code_language: str,
+    tgt_code_language: Optional[str] = None
+):
+    save_dir = os.path.join(save_dir, language, task_type)
+    if tgt_code_language is not None:
+        file_name = f"{language}-{code_language}-to-{tgt_code_language}-triplets.jsonl"
+    else:
+        file_name = f"{language}-{code_language}-triplets.jsonl"
+    save_path = os.path.join(save_dir, file_name)
+    os.makedirs(save_dir, exist_ok=True)
+    return save_path
+
+
+def save_triplets(
+    triplets: list,
+    save_dir: str,
+    task_type: str,
+    language: str,
+    code_language: str,
+    tgt_code_language: Optional[str] = None
+):
+    if len(triplets) == 0:
+        print(f"No triplets to save: {task_type} | {language} | {code_language} | {tgt_code_language}")
+        return
+    
+    save_path = get_save_path(save_dir, task_type, language, code_language, tgt_code_language)
+    query_md5s = set()
+    pos_md5s = set()
+    old_triplets = []
+    if os.path.exists(save_path):
+        with open(save_path, "r", encoding="utf-8") as f:
+            for line in f.readlines():
+                triplet = json.loads(line)
+                old_triplets.append(triplet)
+                query_md5s.add(compute_md5(triplet['query']))
+                pos_md5s.add(compute_md5(triplet['pos'][0]))
+
+    with open(save_path, 'w', encoding='utf-8') as f:
+        for triplet in old_triplets:
+            f.write(json.dumps(triplet, ensure_ascii=False) + '\n')
+        
+        for triplet in triplets:
+            _query_md5 = compute_md5(triplet['query'])
+            _pos_md5 = compute_md5(triplet['pos'][0])
+            if _query_md5 in query_md5s or _pos_md5 in pos_md5s:
+                continue
+            f.write(json.dumps(triplet, ensure_ascii=False) + '\n')
+    print(f"Triplets saved to {save_path}")
+
+
+def main(args):
+    # set seed
+    seed = args.seed
+    if seed is not None:
+        print(f"------------------- Seed set to {seed} -------------------")
+        random.seed(seed)
+    
+    model = args.model
+    model_type = args.model_type
+    port = args.port
+
+    num_samples = args.num_samples
+    
+    task_type = args.task_type
+    language = args.language
+    code_language = args.code_language
+    tgt_code_language = args.tgt_code_language
+
+    corpus_root = args.corpus_root
+    corpus_dir = os.path.join(corpus_root, code_language)
+    doc_length = args.doc_length.split()
+    external_path = args.external_path.split()
+
+    save_dir = args.save_dir
+    cache_dir = args.cache_dir
+    num_processes = min(args.num_processes, int(mp.cpu_count() * 0.8))
+    overwrite = args.overwrite
+    debug_mode = args.debug_mode
+    gen_hard_neg = args.gen_hard_neg
+    
+    save_path = get_save_path(save_dir, task_type, language, code_language, tgt_code_language)
+    # if os.path.exists(save_path) and not overwrite:
+        # data = []
+        # with open(save_path) as f:
+        #     for line in f:
+        #         data.append(json.loads(line))
+        # if len(data) >= num_samples * 0.8:
+        #     print(f"Triplets already exist at {save_path}. Skipping generation.")
+        #     return
+        # else:
+        #     print(f"Triplets already exist at {save_path}. But samples is really small, continue generation.")
+        #     num_samples = int((num_samples - len(data)) * 1.25)  # consider the filtered samples
+
+    corpus_generator = CorpusGenerator(cache_dir)
+
+    examples_dir = args.examples_dir
+    num_examples = args.num_examples
+    if examples_dir is not None:
+        # if task_type in ["single_turn_code_qa", "multi_turn_code_qa"]:
+        #     examples_path = os.path.join(examples_dir, language, task_type, "sample_examples.json")
+        if task_type in ["code_translation_retrieval"]:
+            examples_path = os.path.join(examples_dir, language, task_type,
+                                         f"{code_language}-to-{tgt_code_language}_sample_examples.json")
+        else:
+            examples_path = os.path.join(examples_dir, language, task_type, f"{code_language}_sample_examples.json")
+        try:
+            with open(examples_path, 'r', encoding='utf-8') as f:
+                examples_pool = json.load(f)
+                examples_pool = random.sample(examples_pool,
+                                              min(30, len(examples_pool)))   # sample 30 examples for few-shot generation
+        except:
+            print(f'Error for loading examples from {examples_path}')
+            examples_pool = None
+    else:
+        examples_pool = None
+
+    positives, large_positives = corpus_generator.run(
+        num_samples=num_samples,
+        max_corpus=args.max_corpus,
+        corpus_dir=corpus_dir,
+        doc_length=doc_length,
+        external_path=external_path,
+        source_language=code_language
+    )
+
+    if task_type in ["code_modification_retrieval", "code_comparison_retrieval"]:
+        top1_docs = get_top1([e['text'] for e in positives], args.sim_model_name, [e['text'] for e in large_positives])
+        for i in range(len(top1_docs)):
+            positives[i]['similar'] = top1_docs[i]
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    print("=================== Generate training data ===================")
+    print(f'Task Type: {task_type} | Language: {language} | Code Language: {code_language} | Target Code Language: {tgt_code_language}')
+    start_time = time.time()
+    triplets = gen_triplets(
+        model=model,
+        model_type=model_type,
+        port=port,
+        positives=positives,
+        task_type=task_type,
+        language=language,
+        code_language=code_language,
+        tgt_code_language=tgt_code_language,
+        examples_pool=examples_pool,
+        num_examples=num_examples,
+        thread_count=num_processes,
+        gen_cache_dir=os.path.join(save_dir, language, task_type, "gen_cache_dir"),
+        debug_mode=debug_mode,
+        gen_hard_neg=gen_hard_neg,
+    )
+    save_triplets(
+        triplets=triplets,
+        save_dir=save_dir,
+        task_type=task_type,
+        language=language,
+        code_language=code_language,
+        tgt_code_language=tgt_code_language
+    )
+    end_time = time.time()
+    print("=============================================================")
+    print(f"Time taken: {end_time - start_time:.2f} seconds")
+    print("=============================================================")
+    print("DONE!")
+
+
+if __name__ == "__main__":
+    args = get_args()
+    main(args)

research/BGE_Coder/data_generation/search.py

@@ -0,0 +1,71 @@
+from typing import Optional, List
+
+import faiss
+import numpy as np
+from tqdm import tqdm
+from FlagEmbedding import FlagModel
+
+def create_index(embeddings: np.ndarray, use_gpu: bool = False):
+    index = faiss.IndexFlatIP(len(embeddings[0]))
+    embeddings = np.asarray(embeddings, dtype=np.float32)
+    if use_gpu:
+        co = faiss.GpuMultipleClonerOptions()
+        co.shard = True
+        co.useFloat16 = True
+        index = faiss.index_cpu_to_all_gpus(index, co=co)
+    index.add(embeddings)
+    return index
+
+
+def search(
+        faiss_index: faiss.Index,
+        k: int = 100,
+        query_embeddings: Optional[np.ndarray] = None,
+        load_path: Optional[str] = None
+):
+    if query_embeddings is None:
+        query_embeddings = np.load(load_path)
+
+    query_size = len(query_embeddings)
+
+    all_scores = []
+    all_indices = []
+
+    for i in tqdm(range(0, query_size, 32), desc="Searching"):
+        j = min(i + 32, query_size)
+        query_embedding = query_embeddings[i: j]
+        score, indice = faiss_index.search(query_embedding.astype(np.float32), k=k)
+        all_scores.append(score)
+        all_indices.append(indice)
+
+    all_scores = np.concatenate(all_scores, axis=0)
+    all_indices = np.concatenate(all_indices, axis=0)
+    return all_scores, all_indices
+
+def get_top1(
+        small_docs,
+        encoder_name,
+        docs: List[str],
+        top: int = 1
+):
+    encoder = FlagModel(encoder_name, trust_remote_code=True)
+    doc_emb = encoder.encode_corpus(docs, max_length=512, batch_size=256)
+    small_doc_emb = encoder.encode_corpus(small_docs, max_length=512, batch_size=256)
+    faiss_index = create_index(doc_emb, True)
+    all_scores, all_indices = search(faiss_index, 1000, small_doc_emb)
+    return_docs = []
+    for i in range(len(all_indices)):
+        return_docs.append([])
+        for idx, score in zip(all_indices[i][20:], all_scores[i][20:]):
+            d1 = set(docs[idx].split())
+            d2 = set(small_docs[i].split())
+            if len(d1 & d2) / len(d1 | d2) > 0.95:
+                continue
+            return_docs[-1].append(docs[idx])
+            if len(return_docs[-1]) >= top:
+                break
+        if len(return_docs[-1]) == 0:
+            print(all_indices[i], all_scores[i])
+    # print(return_docs)
+    del faiss_index
+    return return_docs

research/BGE_Coder/data_generation/triplet_generator.py

@@ -0,0 +1,654 @@
+import os
+import json
+import random
+from tqdm import tqdm
+from hashlib import md5
+from warnings import warn
+from typing import List, Optional
+from concurrent.futures import ThreadPoolExecutor
+
+from llm import LLM
+from utils import clean_content
+from constant import TaskType, Task, SPECIAL_TASK_STEPS, \
+    get_task, get_generation_prompt, get_quality_control_prompt, \
+    get_gen_hard_neg_prompt
+
+
+def compute_md5(text: str):
+    return md5(text.encode()).hexdigest()
+
+
+class TripletGenerator(LLM):
+    def __init__(
+        self,
+        model: str = "Qwen2-5-Coder-32B-Instruct",
+        model_type: str = "open-source",
+        port: int = 8000,
+        cache_dir: Optional[str] = None
+    ):
+        super().__init__(model, model_type, port)
+        self.cache_dir = cache_dir
+        if self.cache_dir is not None:
+            os.makedirs(self.cache_dir, exist_ok=True)
+    
+    def _gen_for_code_modification_retrieval(
+        self,
+        task: Task,
+        text: str,
+        text_b: Optional[str] = None,
+        examples: Optional[List[dict]] = None,
+        debug_mode: bool = False,
+        **kwargs
+    ):
+        gen_prompt = get_generation_prompt(
+            task=task,
+            text=text,
+            text_b=text_b,
+            examples=examples,
+            idx=0
+        )
+        response = self.chat(gen_prompt, **kwargs)[0]
+        diff = clean_content(response)
+        gen_prompt = get_generation_prompt(
+            task=task,
+            text=diff,
+            examples=examples,
+            idx=1
+        )
+        response = self.chat(gen_prompt, **kwargs)[0]
+        modification_instr = clean_content(response)
+        
+        query = f"{modification_instr}\n```\n{text}\n```"
+        pos = text_b
+        
+        if debug_mode:
+            result = {
+                "generation_prompt": gen_prompt,
+                "prompt": task.task_instruction,
+                "query": query,
+                "pos": [pos],
+                "neg": []
+            }
+        else:
+            result = {
+                "prompt": task.task_instruction,
+                "query": query,
+                "pos": [pos],
+                "neg": []
+            }
+        return result
+    
+    def _gen_for_code_comparison_retrieval(
+        self,
+        task: Task,
+        text: str,
+        text_b: Optional[str] = None,
+        examples: Optional[List[dict]] = None,
+        debug_mode: bool = False,
+        **kwargs
+    ):
+        gen_prompt = get_generation_prompt(
+            task=task,
+            text=text,
+            text_b=text_b,
+            examples=examples,
+            idx=0
+        )
+        response = self.chat(gen_prompt, **kwargs)[0]
+        diff_question = clean_content(response)
+        query = f"{diff_question}\n\nInput Code:\n```\n{text}\n```\n\nOutput Code:\n```\n{text_b}\n```"
+        gen_prompt = get_generation_prompt(
+            task=task,
+            text=query,
+            examples=examples,
+            idx=1
+        )
+        response = self.chat(gen_prompt, **kwargs)[0]
+        pos = clean_content(response)
+
+        if debug_mode:
+            result = {
+                "generation_prompt": gen_prompt,
+                "prompt": task.task_instruction,
+                "query": query,
+                "pos": [pos],
+                "neg": []
+            }
+        else:
+            result = {
+                "prompt": task.task_instruction,
+                "query": query,
+                "pos": [pos],
+                "neg": []
+            }
+        return result
+    
+    def _gen_for_code_context_retrieval(
+        self,
+        task: Task,
+        text: str,
+        anchor_points: Optional[tuple] = (0.4, 0.7),
+        **kwargs
+    ):
+        former_part, latter_part = self.split_text(
+            text,
+            anchor_points=anchor_points
+        )
+        result = {
+            "prompt": task.task_instruction,
+            "query": former_part,
+            "pos": [latter_part],
+            "neg": []
+        }
+        return result
+    
+    @staticmethod
+    def _arrange_query_and_pos(task: Task, input_text: str, response: str):
+        """
+        Arrange the query and positive example based on the task type.
+        
+        Args:
+        - task: Task
+        - input_text: str
+        - response: str
+        
+        Returns:
+        - query: str
+        - pos: str
+        """
+        # TODO: support more task types, including some special task types.
+        if task.main_task_type in ["text2code", "hybrid"]:
+            query = clean_content(response)
+            pos = input_text
+        else:
+            query = input_text
+            pos = clean_content(response)
+        return query, pos
+    
+    def _gen_for_normal_task(
+        self,
+        task: Task,
+        text: str,
+        examples: Optional[List[dict]] = None,
+        debug_mode: bool = False,
+        **kwargs
+    ):
+        gen_prompt = get_generation_prompt(
+            task=task,
+            text=text,
+            examples=examples
+        )
+        response = self.chat(gen_prompt, **kwargs)[0]
+        
+        # Arrange the query and positive example based on the task type.
+        query, pos = self._arrange_query_and_pos(
+            task=task,
+            input_text=text,
+            response=response
+        )
+        
+        if debug_mode:
+            result = {
+                "generation_prompt": gen_prompt,
+                "prompt": task.task_instruction,
+                "query": query,
+                "pos": [pos],
+                "neg": [],
+                "response": response
+            }
+        else:
+            result = {
+                "prompt": task.task_instruction,
+                "query": query,
+                "pos": [pos],
+                "neg": []
+            }
+        return result
+    
+    def _gen_for_bug_desc_retrieval(
+        self,
+        task: Task,
+        text: str,
+        examples: Optional[List[dict]] = None,
+        debug_mode: bool = False,
+        **kwargs
+    ):
+        gen_prompt = get_generation_prompt(
+            task=task,
+            text=text,
+            examples=examples,
+            idx=0
+        )
+        response = self.chat(gen_prompt, **kwargs)[0]
+        if response is None:
+            raise ValueError("Response is None.")
+        buggy_code = response
+        gen_prompt = get_generation_prompt(
+            task=task,
+            text=buggy_code,
+            examples=examples,
+            idx=1
+        )
+        response = self.chat(gen_prompt, **kwargs)[0]
+        query = clean_content(response)
+        pos = text
+        
+        if debug_mode:
+            result = {
+                "generation_prompt": gen_prompt,
+                "prompt": task.task_instruction,
+                "query": query,
+                "pos": [pos],
+                "neg": []
+            }
+        else:
+            result = {
+                "prompt": task.task_instruction,
+                "query": query,
+                "pos": [pos],
+                "neg": []
+            }
+        return result
+    
+    def _gen_for_two_step_not_use_last(
+        self,
+        task: Task,
+        text: str,
+        examples: Optional[List[dict]] = None,
+        debug_mode: bool = False,
+        reverse_query_pos: bool = False,
+        **kwargs
+    ):
+        gen_prompt = get_generation_prompt(
+            task=task,
+            text=text,
+            idx=0
+        )
+        response = self.chat(gen_prompt, **kwargs)[0]
+        query = clean_content(response)
+        gen_prompt = get_generation_prompt(
+            task=task,
+            text=query,
+            examples=examples,
+            idx=1
+        )
+        response = self.chat(gen_prompt, **kwargs)[0]
+        pos = clean_content(response)
+        if reverse_query_pos:
+            query, pos = pos, query
+
+        if debug_mode:
+            result = {
+                "generation_prompt": gen_prompt,
+                "prompt": task.task_instruction,
+                "query": query,
+                "pos": [pos],
+                "neg": []
+            }
+        else:
+            result = {
+                "prompt": task.task_instruction,
+                "query": query,
+                "pos": [pos],
+                "neg": []
+            }
+        return result
+
+    def _gen_for_two_step_use_last(
+        self,
+        task: Task,
+        text: str,
+        examples: Optional[List[dict]] = None,
+        debug_mode: bool = False,
+        reverse_query_pos: bool = False,
+        **kwargs
+    ):
+        gen_prompt = get_generation_prompt(
+            task=task,
+            text=text,
+            idx=0
+        )
+        response = self.chat(gen_prompt, **kwargs)[0]
+        query = clean_content(response) + f"\n```\n{text}\n```"
+        gen_prompt = get_generation_prompt(
+            task=task,
+            text=query,
+            examples=examples,
+            idx=1
+        )
+        response = self.chat(gen_prompt, **kwargs)[0]
+        pos = clean_content(response)
+        if reverse_query_pos:
+            query, pos = pos, query
+
+        if debug_mode:
+            result = {
+                "generation_prompt": gen_prompt,
+                "prompt": task.task_instruction,
+                "query": query,
+                "pos": [pos],
+                "neg": []
+            }
+        else:
+            result = {
+                "prompt": task.task_instruction,
+                "query": query,
+                "pos": [pos],
+                "neg": []
+            }
+        return result
+
+    def generate_triplets(
+        self,
+        data: dict,
+        task: Task,
+        examples_pool: Optional[List[dict]] = None,
+        num_examples: int = 3,
+        debug_mode: bool = False,
+        **kwargs
+    ):
+        kwargs["remove_thinking"] = not debug_mode
+        
+        result_list = []
+        
+        examples = None
+        if examples_pool is not None:
+            examples = random.sample(examples_pool, min(num_examples, len(examples_pool)))
+
+        try:
+            if task.task_type in SPECIAL_TASK_STEPS:
+                text = data["text"]
+                
+                if task.task_type == TaskType.code_modification_retrieval:
+                    text_b = data["similar"][0]
+                    
+                    result = self._gen_for_code_modification_retrieval(
+                        task=task,
+                        text=text,
+                        text_b=text_b,
+                        examples=examples,
+                        debug_mode=debug_mode
+                    )
+                elif task.task_type == TaskType.code_comparison_retrieval:
+                    text_b = data["similar"][0]
+                    
+                    result = self._gen_for_code_comparison_retrieval(
+                        task=task,
+                        text=text,
+                        text_b=text_b,
+                        examples=examples,
+                        debug_mode=debug_mode
+                    )
+                elif task.task_type == TaskType.bug_desc_retrieval:
+                    result = self._gen_for_bug_desc_retrieval(
+                        task=task,
+                        text=text,
+                        examples=examples,
+                        debug_mode=debug_mode
+                    )
+                elif task.task_type in [
+                    # cf - updated
+                    TaskType.code_issue_discussion_retrieval,
+                    TaskType.code_version_update_retrieval,
+                    TaskType.code_bug_fix_example_retrieval,
+                ]:
+                    result = self._gen_for_two_step_not_use_last(
+                        task=task,
+                        text=text,
+                        examples=examples,
+                        debug_mode=debug_mode,
+                        reverse_query_pos=False
+                    )
+                elif task.task_type in [
+                    # cf - updated
+                    TaskType.code_refactoring_pattern_retrieval,
+                    TaskType.code_style_guideline_example_retrieval,
+                    TaskType.code_migration_retrieval,
+                    # jl - updated
+                    TaskType.code_optimization_hybrid_retrieval,
+                    TaskType.code_best_practices_retrieval,
+                    TaskType.security_vulnerability_fix_retrieval,
+                ]:
+                    result = self._gen_for_two_step_use_last(
+                        task=task,
+                        text=text,
+                        examples=examples,
+                        debug_mode=debug_mode,
+                        reverse_query_pos=False
+                    )
+                else:
+                    raise NotImplementedError(f"Task type {task.task_type} not implemented.")
+            elif task.task_type == TaskType.code_context_retrieval:
+                text = data["text"]
+                
+                result = self._gen_for_code_context_retrieval(
+                    task=task,
+                    text=text,
+                    **kwargs
+                )
+                # NOTE: no need to do quality control for code context retrieval task
+                result_list.append(result)
+                return result_list
+            else:
+                text = data["text"]
+                
+                result = self._gen_for_normal_task(
+                    task=task,
+                    text=text,
+                    examples=examples,
+                    debug_mode=debug_mode,
+                    **kwargs
+                )
+            
+            # print(gen_prompt)
+            # print('================================================')
+            qc_prompt = get_quality_control_prompt(
+                task=task,
+                query=result["query"],
+                pos=result["pos"][0]
+            )
+            # print(qc_prompt)
+            # print('*********************************************************************')
+            response = self.chat(qc_prompt, **kwargs)[0]
+            judge = clean_content(response)
+            # print(response, judge)
+            if "1" in judge:
+                if debug_mode:
+                    result["judge"] = judge
+                    result["judge_response"] = response
+                result_list.append(result)
+            else:
+                if debug_mode:
+                    result["judge"] = judge
+                    result["judge_response"] = response
+                    result_list.append(result)
+        except Exception as e:
+            warn(f"Error: {e}")
+        
+        return result_list
+
+    def gen_hard_negatives(self, result: dict, task: Task, num_negatives: int = 7, **kwargs):
+        gen_hard_neg_prompt = get_gen_hard_neg_prompt(
+            task=task,
+            query=result["query"],
+            pos=result["pos"][0]
+        )
+        response_list = self.chat(gen_hard_neg_prompt, n=num_negatives, **kwargs)
+        for response in response_list:
+            if response is None:
+                continue
+            hard_neg = clean_content(response)
+            result["neg"].append(hard_neg)
+        result["neg"] = list(set(result["neg"]))
+        return result
+
+    def run_single(
+        self,
+        data: dict,
+        task: Task,
+        examples_pool: Optional[List[dict]] = None,
+        num_examples: int = 3,
+        debug_mode: bool = False,
+        gen_hard_neg: bool = False,
+        num_negatives: int = 7,
+        **kwargs
+    ):
+        result_list = []
+
+        docid = compute_md5(data["text"])
+        if self.cache_dir is not None:
+            gen_data_cache_path = os.path.join(self.cache_dir, f"{docid}.json")
+            if os.path.exists(gen_data_cache_path):
+                with open(gen_data_cache_path, "r", encoding="utf-8") as f:
+                    result_list = json.load(f)
+                
+                if len(result_list) > 0:
+                    if gen_hard_neg:
+                        for i in range(len(result_list)):
+                            if len(result_list[i]["neg"]) == 0:
+                                result_list[i] = self.gen_hard_negatives(
+                                    result=result_list[i],
+                                    task=task,
+                                    num_negatives=num_negatives,
+                                    **kwargs
+                                )
+                        # overwrite the cache file
+                        with open(gen_data_cache_path, "w", encoding="utf-8") as f:
+                            json.dump(result_list, f, indent=4, ensure_ascii=False)
+                    return result_list
+
+        triplets = self.generate_triplets(
+            data,
+            task=task,
+            examples_pool=examples_pool,
+            num_examples=num_examples,
+            debug_mode=debug_mode,
+            **kwargs
+        )
+        if len(triplets) == 0:
+            return []
+        
+        result = triplets[0]
+        if debug_mode:
+            result["docid"] = docid
+        
+        if gen_hard_neg:
+            result = self.gen_hard_negatives(
+                result,
+                task=task,
+                num_negatives=num_negatives,
+                **kwargs
+            )
+        
+        result_list.append(result)
+        
+        if self.cache_dir is not None:
+            gen_data_cache_path = os.path.join(self.cache_dir, f"{docid}.json")
+            with open(gen_data_cache_path, "w", encoding="utf-8") as f:
+                json.dump(result_list, f, indent=4, ensure_ascii=False)
+        
+        return result_list
+
+    def run(
+        self,
+        positives: List[dict],
+        task_type: str,
+        language: str = "en",
+        code_language: str = "python",
+        tgt_code_language: Optional[str] = None,
+        examples_pool: Optional[List[dict]] = None,
+        num_examples: int = 3,
+        tqdm_desc: str = "Generating triplets",
+        debug_mode: bool = False,
+        gen_hard_neg: bool = False,
+        num_negatives: int = 7,
+        thread_count: int = 1,
+        **kwargs
+    ):
+        task = get_task(
+            task_type=task_type,
+            language=language,
+            code_language=code_language,
+            tgt_code_language=tgt_code_language
+        )
+        
+        result_list = []
+
+        def process_positive(positive):
+            return self.run_single(
+                data=positive,
+                task=task,
+                examples_pool=examples_pool,
+                num_examples=num_examples,
+                debug_mode=debug_mode,
+                gen_hard_neg=gen_hard_neg,
+                num_negatives=num_negatives,
+                **kwargs
+            )
+        # Use thread pool for parallel processing with tqdm progress bar.
+        with ThreadPoolExecutor(max_workers=thread_count) as executor:
+            results = list(tqdm(executor.map(
+                process_positive,
+                positives
+            ), total=len(positives), desc=tqdm_desc))
+
+        # Collect results into result_list.
+        for res in results:
+            if isinstance(res, list):
+                result_list.extend(res)
+            else:
+                result_list.append(res)
+        # result_list.extend(results)
+
+        return result_list
+
+    def run_for_gen_neg(
+        self,
+        pairs: List[dict],
+        task_type: str,
+        language: str = "en",
+        code_language: str = "python",
+        tgt_code_language: Optional[str] = None,
+        examples_pool: Optional[List[dict]] = None,
+        num_examples: int = 3,
+        tqdm_desc: str = "Generating triplets",
+        debug_mode: bool = False,
+        gen_hard_neg: bool = False,
+        num_negatives: int = 7,
+        thread_count: int = 1,
+        **kwargs
+    ):
+        task = get_task(
+            task_type=task_type,
+            language=language,
+            code_language=code_language,
+            tgt_code_language=tgt_code_language
+        )
+        
+        result_list = []
+
+        def gen_single_negative(pair):
+            result = self.gen_hard_negatives(
+                pair,
+                task=task,
+                num_negatives=num_negatives,
+                **kwargs
+            )
+            return [result]
+
+        # Use thread pool for parallel processing with tqdm progress bar.
+        with ThreadPoolExecutor(max_workers=thread_count) as executor:
+            results = list(tqdm(executor.map(
+                gen_single_negative,
+                pairs
+            ), total=len(pairs), desc=tqdm_desc))
+
+        # Collect results into result_list.
+        for res in results:
+            if isinstance(res, list):
+                result_list.extend(res)
+            else:
+                result_list.append(res)
+        # result_list.extend(results)
+
+        return result_list

research/BGE_Coder/data_generation/utils.py

@@ -0,0 +1,128 @@
+import re
+
+
+def clean_content(content: str):
+    if content is None:
+        raise ValueError("content is None.")
+    
+    content = content.split('</think>')[-1].strip('\n').strip()
+    
+    if content.startswith('\"') and content.endswith('\"'):
+        content = content[1:-1]
+    
+    if content.startswith("```\n") and content.endswith("\n```"):
+        content = content[4:-4]
+    
+    return content
+
+
+def clean_code(code: str, lang: str, length_threshold: int = 30) -> str:
+    cleaned_code = code.strip('\ufeff').strip()
+    if not cleaned_code:
+        return ''
+
+    def clean_empty_lines(text: str) -> str:
+        return re.sub(r'\n\s*\n', '\n', text).strip()
+
+    # 各语言函数/类定义检测正则表达式
+    function_patterns = {
+        "java": r"(?m)^(?!\s*(import|package)\b).*\b(public\s+class|class\s+\w+|void\s+main|new\s+\w+\(|@Override)\b",
+        "python": r"(?m)^(?!\s*(import|from\s+\S+\s+import)\b).*\b(def\s+\w+|class\s+\w+|=\s*\S+|if\s+[:\w]|print\s+)",
+        "javascript": r"(?m)^(?!\s*(import|require\(|export\s)).*\b(function\s+\w+|const\s+\w+|=>|\(\)\s*=>|console\.log)",
+        "php": r"(?m)^(?!\s*(include|require|use)\b).*\b(function\s+\w+|echo\s+\S+|class\s+\w+)",
+        "ruby": r"(?m)^(?!\s*(require|load)\b).*\b(class\s+\w+|def\s+\w+|puts\s+\S+)",
+        "go": r"(?m)^(?!\s*import\b).*\bfunc\s+main\s*\(|type\s+\w+\s+struct",
+        "c#": r"(?m)^(?!\s*using\b).*\b(class\s+\w+|void\s+Main\s*\()",
+        "cplusplus": r"(?m)^(?!#include\b).*\b(int\s+main\s*\(|class\s+\w+|void\s+\w+\s*\(.*\)\s*{)",
+        "c": r"(?m)^(?!#include\b).*\b(int\s+main\s*\(|void\s+\w+\s*\(.*\)\s*{)",
+        "rust": r"(?m)^(?!\s*use\b).*\b(fn\s+main\s*\(|struct\s+\w+|impl\s+\w+)",
+        "typescript": r"(?m)^(?!\s*(import|require\(|export\s)).*\b(interface\s+\w+|class\s+\w+|function\s+\w+)",
+        "perl": r"(?m)^(?!\s*(use|require)\b).*\b(sub\s+\w+|my\s+\$\w+|print\s+\S+)",
+        "shell": r"(?m)^(?!\s*(source|\.)\s).*\b(function\s+\w+|if\s+\[|\$\(|echo\s+\S+)",
+        "sql": r"(?i)\b(CREATE\s+TABLE|SELECT\s+\*|INSERT\s+INTO|UPDATE\s+\w+|DELETE\s+FROM)\b",
+        "batchfile": r"(?m)^(?!\s*@?call\b).*\b(echo\s+\S+|set\s+\w+|if\s+.*\s+==\s+)",
+        "fortran": r"(?mi)^(?!\s*use\b).*\b(program\s+\w+|subroutine\s+\w+|do\s+\d+\s*,\s*\d+)",
+        "haskell": r"(?m)^(?!\s*import\b).*\b(main\s*=\s*do|data\s+\w+|putStrLn\s+\S+)",
+        "lua": r"(?m)^(?!\s*require\b).*\b(function\s+\w+|local\s+\w+|print\s*\()",
+        "powershell": r"(?m)^(?!\s*Import-Module\b).*\b(function\s+\w+|Write-Host\s+\S+|\$\w+\s*=)",
+        "visual_basic": r"(?m)^(?!\s*Imports\b).*\b(Module\s+\w+|Sub\s+Main|Class\s+\w+)"
+    }
+
+    # 各语言注释处理规则
+    comment_patterns = {
+        'java': (r'//.*?$|/\*.*?\*/|\\/\\/.*?$|\\/\*.*?\*\\/', re.DOTALL | re.MULTILINE),
+        'python': (r'#.*?$', re.MULTILINE),
+        'javascript': (r'//.*?$|/\*.*?\*/|\\/\\/.*?$|\\/\*.*?\*\\/', re.DOTALL | re.MULTILINE),
+        'php': (r'//.*?$|#.*?$|/\*.*?\*/|\\/\\/.*?$|#.*?$|\\/\*.*?\*\\/', re.DOTALL | re.MULTILINE),
+        'ruby': (r'#.*', re.MULTILINE),
+        'go': (r'//.*?$|/\*.*?\*/|\\/\\/.*?$|\\/\*.*?\*\\/', re.DOTALL | re.MULTILINE),
+        'csharp': (r'//.*?$|/\*.*?\*/|\\/\\/.*?$|\\/\*.*?\*\\/', re.DOTALL | re.MULTILINE),
+        'cplusplus': (r'//.*?$|/\*.*?\*/|\\/\\/.*?$|\\/\*.*?\*\\/', re.DOTALL | re.MULTILINE),
+        'c': (r'//.*?$|/\*.*?\*/|\\/\\/.*?$|\\/\*.*?\*\\/', re.DOTALL | re.MULTILINE),
+        'rust': (r'//.*?$|/\*.*?\*/|\\/\\/.*?$|\\/\*.*?\*\\/', re.DOTALL | re.MULTILINE),
+        'typescript': (r'//.*?$|/\*.*?\*/|\\/\\/.*?$|\\/\*.*?\*\\/', re.DOTALL | re.MULTILINE),
+        'perl': (r'#.*', re.MULTILINE),
+        'shell': (r'#.*', re.MULTILINE),
+        'sql': (r'--.*?$|/\*.*?\*/', re.DOTALL),
+        'batchfile': (r'^\s*(REM|@REM|::).*', re.MULTILINE | re.IGNORECASE),
+        'fortran': (r'!.*', re.MULTILINE),
+        'haskell': (r'--.*', re.MULTILINE),
+        'lua': (r'--.*?$|--\[\[.*?\]\]', re.DOTALL),
+        'powershell': (r'<#.*?#>|#.*', re.DOTALL),
+        'visual_basic': (r"'.*", re.MULTILINE),
+    }
+
+    # 执行注释清理
+    if lang in comment_patterns:
+        pattern, flags = comment_patterns[lang]
+        cleaned_code = re.sub(pattern, '', cleaned_code, flags=flags)
+        cleaned_code = clean_empty_lines(cleaned_code)
+
+    # 特殊语言处理规则
+    if lang == 'fortran':
+        cleaned_code = re.sub(r'^[Cc*].*', '', cleaned_code, flags=re.MULTILINE)
+    elif lang == 'sql':
+        cleaned_code = re.sub(r'/\*.*?\*/', '', cleaned_code, flags=re.DOTALL)
+    elif lang == 'python':
+        cleaned_code = re.sub(r'^\s*#.*', '', cleaned_code, flags=re.MULTILINE)
+
+    # 函数定义检测及内容验证
+    def has_valid_code(text: str, lang: str) -> bool:
+        pattern = function_patterns.get(lang)
+        if not pattern:
+            return len(text.strip()) > 0
+        
+        # 增强检测逻辑
+        if lang == 'batchfile':
+            return bool(re.search(r'^\s*@?echo\b|:\w+', text, re.MULTILINE))
+        if lang == 'shell':
+            return bool(re.search(r'^\s*(if|for|while|case|echo|export|shopt|source)\b', text, re.MULTILINE))
+        if lang == 'python':
+            if re.search(r'^\s*(def|class)\s+\w+', text, re.MULTILINE):
+                return bool(re.search(r'^\s+[^\s#]', text, re.MULTILINE))
+            return False
+        if lang == 'ruby':
+            return bool(re.search(r'(def\s+\w+|class\s+\w+).*?\n\s+[^\s#]', text, re.MULTILINE))
+        return bool(re.search(pattern, text, re.DOTALL | re.MULTILINE))
+
+    # 最终有效性检查
+    if not has_valid_code(cleaned_code, lang):
+        return ''
+    
+    cleaned_code = cleaned_code.strip('\ufeff').strip()
+    
+    if len(cleaned_code) < length_threshold:
+        return ''
+
+    return cleaned_code
+
+
+if __name__ == "__main__":
+    test_text = "\/\/ ----------------------------------------------------------------------\n\/\/ ----------------------------------------------------------------------\n\/\/\n\/\/ File:      StrMaxProjection.h\n\/\/ Author:    mgrosso \n\/\/ Created:   Mon Jul 17 14:39:22 PDT 2006 on caliban\n\/\/ Project:   \n\/\/ Purpose:   \n\/\/ \n\/\/ $Id$\n\/\/ ----------------------------------------------------------------------\n\/\/ ----------------------------------------------------------------------\n\n#ifndef STRMAXPROJECTION_H\n#define STRMAXPROJECTION_H 1\n\n#include \"StrMinProjection.h\"\n\nclass StrMaxProjection : public StrMinProjection\n{\n    public:\n    StrMaxProjection(ExpressionPtr &operand);\n    virtual ~StrMaxProjection();\n    virtual     AbstractProjectionPtr        copy();\n\n    protected:\n    int compare(const char *lhs, const char *rhs);\n\n    private:\n    \/\/not implemented\n    StrMaxProjection();\n    StrMaxProjection( const StrMaxProjection &rhs );\n    StrMaxProjection &operator=( const StrMaxProjection &rhs );\n};\n\n#endif \/* STRMAXPROJECTION_H *\/"
+    
+    result = clean_code(test_text, "c", 200)
+    print(result)
+    
+    test_text = "\/**\n * Copyright (c) Microsoft Corporation. All rights reserved.\n * Licensed under the MIT License. See License.txt in the project root for\n * license information.\n *\n * Code generated by Microsoft (R) AutoRest Code Generator.\n *\/\n\npackage com.microsoft.azure.management.datafactory.v2018_06_01;\n\nimport com.fasterxml.jackson.annotation.JsonProperty;\nimport com.fasterxml.jackson.annotation.JsonTypeInfo;\nimport com.fasterxml.jackson.annotation.JsonTypeName;\n\n\/**\n * The location of Google Cloud Storage dataset.\n *\/\n@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, include = JsonTypeInfo.As.PROPERTY, property = \"type\", defaultImpl = GoogleCloudStorageLocation.class)\n@JsonTypeName(\"GoogleCloudStorageLocation\")\npublic class GoogleCloudStorageLocation extends DatasetLocation {\n    \/**\n     * Specify the bucketName of Google Cloud Storage. Type: string (or\n     * Expression with resultType string).\n     *\/\n    @JsonProperty(value = \"bucketName\")\n    private Object bucketName;\n\n    \/**\n     * Specify the version of Google Cloud Storage. Type: string (or Expression\n     * with resultType string).\n     *\/\n    @JsonProperty(value = \"version\")\n    private Object version;\n\n    \/**\n     * Get specify the bucketName of Google Cloud Storage. Type: string (or Expression with resultType string).\n     *\n     * @return the bucketName value\n     *\/\n    public Object bucketName() {\n        return this.bucketName;\n    }\n\n    \/**\n     * Set specify the bucketName of Google Cloud Storage. Type: string (or Expression with resultType string).\n     *\n     * @param bucketName the bucketName value to set\n     * @return the GoogleCloudStorageLocation object itself.\n     *\/\n    public GoogleCloudStorageLocation withBucketName(Object bucketName) {\n        this.bucketName = bucketName;\n        return this;\n    }\n\n    \/**\n     * Get specify the version of Google Cloud Storage. Type: string (or Expression with resultType string).\n     *\n     * @return the version value\n     *\/\n    public Object version() {\n        return this.version;\n    }\n\n    \/**\n     * Set specify the version of Google Cloud Storage. Type: string (or Expression with resultType string).\n     *\n     * @param version the version value to set\n     * @return the GoogleCloudStorageLocation object itself.\n     *\/\n    public GoogleCloudStorageLocation withVersion(Object version) {\n        this.version = version;\n        return this;\n    }\n\n}"
+    result = clean_code(test_text, "java", 200)
+    print(result)

research/BGE_Coder/evaluation/coderag_eval/eval.sh

@@ -0,0 +1,22 @@
+cd ./code-rag-bench/retrieval/
+
+output_dir='result'
+
+for dataset_name in "humaneval" "mbpp" "repoeval" "ds1000_all_completion" "odex_en" "swe-bench-lite"
+do
+echo "dataset_name: ${dataset_name}"
+python main.py \
+    --embedder_name_or_path BAAI/bge-code-v1 \
+    --embedder_model_class decoder-only-base \
+    --query_instruction_format_for_retrieval '<instruct>{}\n<query>{}' \
+    --embedder_query_max_length 2048 \
+    --embedder_passage_max_length 2048 \
+    --trust_remote_code True \
+    --pooling_method last_token \
+    --embedder_batch_size 64 \
+    --devices cuda:0 cuda:1 cuda:2 cuda:3 cuda:4 cuda:5 cuda:6 cuda:7 \
+    --cache_dir ./cache \
+    --dataset $dataset_name \
+    --output_file ../../${output_dir}/${dataset_name}_output.json \
+    --results_file ../../${output_dir}/${dataset_name}_results.json
+done

research/BGE_Coder/evaluation/coderag_eval/prepare_data.sh

@@ -0,0 +1,7 @@
+cd ./code-rag-bench/retrieval/
+
+for dataset_name in "humaneval" "mbpp" "live_code_bench" "ds1000" "odex" "repoeval_repo" "swebench_repo"
+do
+echo "dataset_name: ${dataset_name}"
+PYTHONPATH=./ python create/${dataset_name}.py
+done

research/BGE_Coder/evaluation/coderag_eval/test/arguments.py

@@ -0,0 +1,34 @@
+from typing import List
+from dataclasses import dataclass, field
+
+from FlagEmbedding.abc.evaluation import (
+    AbsEvalModelArgs as CodeRAGEvalModelArgs,
+)
+
+@dataclass
+class CodeRAGEvalArgs:
+    dataset: str = field(
+        default='humaneval',
+        metadata={
+            "help": "Task to evaluate. Default: humaneval. Available tasks: "
+                    "['humaneval', 'mbpp', 'live_code_bench', 'ds1000', 'odex', 'repoeval_repo', 'swebench_repo', 'code_search_net']",
+        }
+    )
+    max_length: int = field(
+        default=2048, metadata={"help": "Max length to use for evaluation."}
+    )
+    batch_size: int = field(
+        default=64, metadata={"help": "Batch size for evaluation."}
+    )
+    output_file: str = field(
+        default="outputs.json",
+        metadata={
+            "help": "Specify the filepath if you want to save the retrieval (evaluation) results."
+        },
+    )
+    results_file: str = field(
+        default="results.json",
+        metadata={
+            "help": "Specify the filepath if you want to save the retrieval results."
+        },
+    )

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)

research/BGE_Coder/evaluation/coderag_eval/test/main.py

@@ -0,0 +1,313 @@
+import os
+import json
+import random
+import logging
+import pathlib
+import argparse
+import numpy as np
+from time import time
+from datasets import load_dataset
+from beir import util, LoggingHandler
+from beir.retrieval import models
+from beir.datasets.data_loader import GenericDataLoader
+from beir.retrieval.evaluation import EvaluateRetrieval
+from beir.retrieval.search.dense import DenseRetrievalExactSearch as DRES
+from tqdm import tqdm
+from transformers import HfArgumentParser
+
+from arguments import CodeRAGEvalArgs, CodeRAGEvalModelArgs
+from prompts import get_task_def_by_task_name
+from FlagEmbedding import FlagLLMModel, FlagModel
+
+
+def get_model(model_args: CodeRAGEvalModelArgs):
+    embedder_name_or_path = model_args.embedder_name_or_path
+
+    if model_args.embedder_model_class == "encoder-only-base":
+        embedder = FlagModel(
+            model_name_or_path=embedder_name_or_path,
+            normalize_embeddings=model_args.normalize_embeddings,
+            pooling_method=model_args.pooling_method,
+            use_fp16=model_args.use_fp16,
+            query_instruction_for_retrieval=model_args.query_instruction_for_retrieval,
+            query_instruction_format=model_args.query_instruction_format_for_retrieval,
+            devices=model_args.devices,
+            trust_remote_code=model_args.trust_remote_code,
+            cache_dir=model_args.cache_dir,
+            batch_size=model_args.embedder_batch_size,
+            query_max_length=model_args.embedder_query_max_length,
+            passage_max_length=model_args.embedder_passage_max_length,
+        )
+    elif model_args.embedder_model_class == "decoder-only-base":
+        embedder = FlagLLMModel(
+            model_name_or_path=embedder_name_or_path,
+            normalize_embeddings=model_args.normalize_embeddings,
+            pooling_method=model_args.pooling_method,
+            use_fp16=model_args.use_fp16,
+            query_instruction_for_retrieval=model_args.query_instruction_for_retrieval,
+            query_instruction_format=model_args.query_instruction_format_for_retrieval,
+            devices=model_args.devices,
+            examples_for_task=model_args.examples_for_task,
+            examples_instruction_format=model_args.examples_instruction_format,
+            trust_remote_code=model_args.trust_remote_code,
+            cache_dir=model_args.cache_dir,
+            batch_size=model_args.embedder_batch_size,
+            query_max_length=model_args.embedder_query_max_length,
+            passage_max_length=model_args.embedder_passage_max_length,
+        )
+    else:
+        raise ValueError(f"Invalid model class: {model_args.embedder_model_class}")
+    embedder.model.config._name_or_path = model_args.embedder_name_or_path
+
+    class CustomFlagModel:
+        def __init__(self, model):
+            self.model = model
+
+        def encode_queries(self, queries, show_progress_bar, convert_to_tensor, **kwargs):
+            if isinstance(queries, str):
+                queries = [queries]
+
+            if isinstance(queries[0], dict):
+                queries = [(e.get('title') + ' ' + e['text']).strip() for e in queries]
+
+            return self.model.encode_queries(queries, **kwargs)
+
+        def encode_corpus(self, corpus, show_progress_bar, convert_to_tensor, **kwargs):
+            if isinstance(corpus, str):
+                corpus = [corpus]
+
+            if isinstance(corpus[0], dict):
+                corpus = [(e.get('title') + ' ' + e['text']).strip() for e in corpus]
+
+            return self.model.encode_corpus(corpus, **kwargs)
+
+        def encode(self, corpus, show_progress_bar, convert_to_tensor, **kwargs):
+            if isinstance(corpus, str):
+                corpus = [corpus]
+
+            if isinstance(corpus[0], dict):
+                corpus = [(e.get('title') + ' ' + e['text']).strip() for e in corpus]
+
+            return self.model.encode(corpus, **kwargs)
+
+    return CustomFlagModel(embedder)
+
+#### Just some code to print debug information to stdout
+logging.basicConfig(format='%(asctime)s - %(message)s',
+                    datefmt='%Y-%m-%d %H:%M:%S',
+                    level=logging.INFO,
+                    handlers=[LoggingHandler()])
+
+
+def get_top_docs(results: dict, corpus: dict, task_id: str, topk: int = 10) -> list[str]:
+    if task_id not in results: return []
+    doc_scores = results[task_id]
+    doc_scores_sorted = sorted(doc_scores.items(), key=lambda item: item[1], reverse=True)
+    doc_scores_sorted = doc_scores_sorted[:topk]
+    doc_code_snippets = [corpus[code_id] for code_id, score in doc_scores_sorted]
+    return doc_code_snippets
+
+
+def main(
+    eval_args: CodeRAGEvalArgs,
+    model_args: CodeRAGEvalModelArgs
+):
+    args = eval_args
+
+    embedder = get_model(model_args)
+    model = DRES(
+        embedder,
+        batch_size=args.batch_size,
+        corpus_chunk_size=512 * 9999
+    )
+    retriever = EvaluateRetrieval(model, score_function="dot")
+
+    if args.dataset.startswith("swe-bench") or args.dataset.startswith("repoeval"):
+        all_eval_results = []
+
+        if args.dataset.startswith("swe-bench"):
+            swebench = load_dataset("princeton-nlp/SWE-bench_Lite")["test"]
+            all_top_docs = [[] for _ in swebench]
+
+        instance_list = [i for i in os.listdir("datasets") if i.startswith(f"{args.dataset}_")]
+        instance_list_filtered = []
+
+        for ins_dir in tqdm(instance_list):
+            logging.info("Instance Repo: {}".format(ins_dir))
+            # load data and perform retrieval
+            corpus, queries, qrels = GenericDataLoader(
+                data_folder=os.path.join("datasets", ins_dir)
+            ).load(split="test")
+            logging.info(f"Instance #{ins_dir}: #{len(corpus)} corpus, #{len(queries)} queries")
+
+            start_time = time()
+            if len(queries) == 1:
+                queries.update({"dummy": "dummy"})
+            results = retriever.retrieve(corpus, queries)
+            if "dummy" in queries:
+                queries.pop("dummy")
+                results.pop("dummy")
+            end_time = time()
+            logging.info("Time taken to retrieve: {:.2f} seconds".format(end_time - start_time))
+
+            # get topk retrieved docs
+            if args.dataset.startswith("swe-bench"):
+                indices = [i for i, ex in enumerate(swebench) if ex["instance_id"] in queries]
+                for index in indices:
+                    instance_id = swebench[index]["instance_id"]
+                    all_top_docs[index] = get_top_docs(results, corpus, instance_id)
+            elif args.dataset.startswith("repoeval"):
+                args.dataset_path = "output/repoeval/datasets/function_level_completion_2k_context_codex.test.clean.jsonl"
+                tasks = [json.loads(line.strip()) for line in open(args.dataset_path, 'r')]
+                prompts, references, docs, metadatas = [], [], [], []
+                for task in tasks:
+                    if task["metadata"]["task_id"] not in queries: continue
+                    prompts.append(task["prompt"])  # save full prompt
+                    references.append(task["metadata"]["ground_truth"])
+                    docs.append(get_top_docs(
+                        results=results, corpus=corpus, task_id=task["metadata"]["task_id"],
+                    ))
+                    metadatas.append(task["metadata"])
+                assert len(prompts) == len(references) == len(docs)
+                dataset = [
+                    {"prompt": p, "reference": r, "docs": d, "metadata": m}
+                    for p, r, d, m in zip(prompts, references, docs, metadatas)
+                ]
+                with open(args.results_file, "a") as fout:
+                    for curr in dataset:
+                        fout.write(json.dumps(curr) + "\n")
+            else:
+                raise ValueError(f"`dataset` should starts with either 'swe-bench' or 'repoeval'.")
+
+            # evaluate retrieval results
+            if len(qrels) == 0:
+                logging.info("No qrels found for this dataset.")
+                return
+            logging.info("Retriever evaluation for k in: {}".format(retriever.k_values))
+            ndcg, _map, recall, precision = retriever.evaluate(qrels, results, retriever.k_values)
+            mrr = retriever.evaluate_custom(qrels, results, retriever.k_values, metric="mrr")
+            eval_results = {
+                "ndcg": ndcg, "mrr": mrr,
+                "recall": recall, "precision": precision,
+                "time": end_time - start_time
+            }
+            logging.info(f"Instance #{ins_dir}: {eval_results}")
+            all_eval_results.append(eval_results)
+
+            with open(args.output_file + "_all", "w") as f:
+                json.dump(all_eval_results, f)
+
+        if args.dataset.startswith("swe-bench"):
+            swebench = swebench.add_column("docs", all_top_docs)
+            swebench.to_json(args.results_file)
+
+        avg_eval_results = {}
+        for k, v_dict in all_eval_results[0].items():
+            if isinstance(v_dict, dict):
+                avg_v_dict = {}
+                for vk, vv in v_dict.items():
+                    avg_vv = sum([e[k][vk] for e in all_eval_results]) / len(all_eval_results)
+                    avg_v_dict[vk] = avg_vv
+                avg_eval_results.update(avg_v_dict)
+            elif isinstance(v_dict, float):
+                avg_v = sum([e[k] for e in all_eval_results]) / len(all_eval_results)
+                avg_eval_results[k] = avg_v
+            else:
+                raise ValueError
+        print("Average Eval Results: ", avg_eval_results)
+        with open(args.output_file, "w") as f:
+            json.dump(avg_eval_results, f)
+    else:
+        dataset = args.dataset
+        corpus, queries, qrels = GenericDataLoader(data_folder=os.path.join("datasets", args.dataset)).load(
+            split="test")
+        #### Retrieve dense results (format of results is identical to qrels)
+        start_time = time()
+        results = retriever.retrieve(corpus, queries)
+        end_time = time()
+        print("Time taken to retrieve: {:.2f} seconds".format(end_time - start_time))
+
+        if args.dataset in ["humaneval", "mbpp", "apps"]:
+            if args.dataset == "humaneval":
+                ds = load_dataset("openai_humaneval")
+                id_key = "task_id"
+            elif args.dataset == "mbpp":
+                ds = load_dataset("mbpp")
+                id_key = "task_id"
+            elif args.dataset == "apps":
+                ds = load_dataset("codeparrot/apps")
+                id_key = "problem_id"
+            all_top_docs = []
+            for task_id in ds["test"][id_key]:
+                all_top_docs.append(get_top_docs(results, corpus, f"{task_id}_doc"))
+            ds["test"] = ds["test"].add_column("docs", all_top_docs)
+            ds["test"].to_json(args.results_file)  # this outputs to arrow format and read as .jsonl
+        elif args.dataset.startswith("odex"):
+            lang = args.dataset.split("_")[-1]
+            ds = load_dataset("neulab/odex", lang, trust_remote_code=True)
+            all_top_docs = []
+            for idx, task_id in enumerate(ds["test"]["task_id"]):
+                all_top_docs.append(get_top_docs(results, corpus, f"{idx}_{task_id}"))
+            ds["test"] = ds["test"].add_column("docs", all_top_docs)
+            ds["test"].to_json(args.results_file)  # this outputs to arrow format and read as .jsonl
+        elif args.dataset.startswith("ds1000"):
+            _, key, mode = args.dataset.split("_")
+            key = key.capitalize()
+            mode = mode.capitalize()
+            from create.ds1000 import get_dataset
+            source_dir = pathlib.Path(__file__).parent / "ds"
+            data = get_dataset(source_dir, mode=mode, key=key)
+            all_docs = []
+            example_ids = []
+            for item in data:
+                example = item.data
+                example_id = f"{example['lib']}_{example['perturbation_origin_id']}"
+                all_docs.append(get_top_docs(results, corpus, example_id))
+                example_ids.append(example_id)
+            assert len(all_docs) == len(
+                example_ids), f"length of all_docs should be {len(example_ids)}, now is {len(all_docs)}"
+            with open(args.results_file, "w+") as fout:
+                for idx, all_doc in enumerate(all_docs):
+                    fout.write(json.dumps({"example_id": example_id,
+                                           "docs": all_doc}) + "\n")
+        else:
+            with open(args.results_file, 'w+') as fw:
+                for curr in results:
+                    fw.write(json.dumps({curr: results[curr]}) + "\n")
+
+        #### Evaluate your retrieval using NDCG@k, MAP@K ...
+        if len(qrels) == 0:
+            logging.info("No qrels found for this dataset.")
+            return
+        logging.info("Retriever evaluation for k in: {}".format(retriever.k_values))
+        ndcg, _map, recall, precision = retriever.evaluate(qrels, results, retriever.k_values)
+
+        mrr = retriever.evaluate_custom(qrels, results, retriever.k_values, metric="mrr")
+        recall_cap = retriever.evaluate_custom(qrels, results, retriever.k_values, metric="r_cap")
+        hole = retriever.evaluate_custom(qrels, results, retriever.k_values, metric="hole")
+
+        all_results = {"ndcg": ndcg, "mrr": mrr, "recall": recall, "precision": precision,
+                       "time": end_time - start_time}
+        with open(args.output_file, "w") as f:
+            json.dump(all_results, f)
+        #### Print top-k documents retrieved ####
+        top_k = 3
+
+        query_id, ranking_scores = random.choice(list(results.items()))
+        scores_sorted = sorted(ranking_scores.items(), key=lambda item: item[1], reverse=True)
+        logging.info("Query : %s\n" % queries[query_id])
+
+        for rank in range(top_k):
+            doc_id = scores_sorted[rank][0]
+            # Format: Rank x: ID [Title] Body
+            logging.info(
+                "Rank %d: %s [%s] - %s\n" % (rank + 1, doc_id, corpus[doc_id].get("title"), corpus[doc_id].get("text")))
+
+
+if __name__ == "__main__":
+    parser = HfArgumentParser((
+        CodeRAGEvalArgs,
+        CodeRAGEvalModelArgs
+    ))
+    eval_args, model_args = parser.parse_args_into_dataclasses()
+    main(eval_args, model_args)

research/BGE_Coder/evaluation/coderag_eval/test/prompts.py

@@ -0,0 +1,18 @@
+from typing import Dict
+
+
+def get_task_def_by_task_name(task_name: str) -> str:
+    task_name_to_instruct: Dict[str, str] = {
+        'humaneval': 'Given a question that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.',
+        'mbpp': 'Given a textual explanation of code functionality, retrieve the corresponding code implementation.',
+        'ds1000_all_completion': 'Given a question that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.',
+        'odex_en': 'Given a question, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.',
+        'odex_es': 'Given a question, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.',
+        'odex_ja': 'Given a question, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.',
+        'odex_ru': 'Given a question, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.',
+        'repoeval': 'Given a code snippet and a new function name, retrieve the implementation of the function.',
+        # 'repoeval': 'Given a piece of code segment, retrieve the code segment that is the latter part of the code.',
+        'swe-bench-lite': 'Given a code snippet containing a bug and a natural language description of the bug or error, retrieve code snippets that demonstrate solutions or fixes for similar bugs or errors (the desired documents).'
+    }
+    
+    return task_name_to_instruct[task_name]

research/BGE_Coder/evaluation/coir_eval/arguments.py

@@ -0,0 +1,69 @@
+from typing import List
+from dataclasses import dataclass, field
+
+from FlagEmbedding.abc.evaluation import (
+    AbsEvalModelArgs as COIREvalModelArgs,
+)
+
+
+def coir_tasks():
+    return [
+        "apps",
+        "codefeedback-mt",
+        "codefeedback-st",
+        "CodeSearchNet-ccr-go",
+        "CodeSearchNet-ccr-java",
+        "CodeSearchNet-ccr-javascript",
+        "CodeSearchNet-ccr-php",
+        "CodeSearchNet-ccr-python",
+        "CodeSearchNet-ccr-ruby",
+        "CodeSearchNet-go",
+        "CodeSearchNet-java",
+        "CodeSearchNet-javascript",
+        "CodeSearchNet-php",
+        "CodeSearchNet-python",
+        "CodeSearchNet-ruby",
+        "codetrans-contest",
+        "codetrans-dl",
+        "cosqa",
+        "stackoverflow-qa",
+        "synthetic-text2sql"
+    ]
+
+
+@dataclass
+class COIREvalArgs:
+    output_dir: str = field(
+        default="./results", metadata={"help": "Path to save results."}
+    )
+    tasks: List[str] = field(
+        default_factory=coir_tasks,
+        metadata={
+            "help": "Tasks to evaluate. Default: None. Available tasks: ['apps', 'codefeedback-mt', 'codefeedback-st', 'CodeSearchNet-ccr-go', 'CodeSearchNet-ccr-java', 'CodeSearchNet-ccr-javascript', 'CodeSearchNet-ccr-php', 'CodeSearchNet-ccr-python', 'CodeSearchNet-ccr-ruby', 'CodeSearchNet-go', 'CodeSearchNet-java', 'CodeSearchNet-javascript', 'CodeSearchNet-php', 'CodeSearchNet-python', 'CodeSearchNet-ruby', 'codetrans-contest', 'codetrans-dl', 'cosqa', 'stackoverflow-qa', 'synthetic-text2sql']",
+            "choices": [
+                "apps",
+                "codefeedback-mt",
+                "codefeedback-st",
+                "CodeSearchNet-ccr-go",
+                "CodeSearchNet-ccr-java",
+                "CodeSearchNet-ccr-javascript",
+                "CodeSearchNet-ccr-php",
+                "CodeSearchNet-ccr-python",
+                "CodeSearchNet-ccr-ruby",
+                "CodeSearchNet-go",
+                "CodeSearchNet-java",
+                "CodeSearchNet-javascript",
+                "CodeSearchNet-php",
+                "CodeSearchNet-python",
+                "CodeSearchNet-ruby",
+                "codetrans-contest",
+                "codetrans-dl",
+                "cosqa",
+                "stackoverflow-qa",
+                "synthetic-text2sql"
+            ]
+        }
+    )
+    use_special_instructions: bool = field(
+        default=False, metadata={"help": "Whether to use specific instructions in `prompts.py` for evaluation. Default: False"}
+    )

research/BGE_Coder/evaluation/coir_eval/eval.sh

@@ -0,0 +1,16 @@
+output_dir=result
+
+python main.py \
+    --output_dir ${output_dir} \
+    --use_special_instructions True \
+    --embedder_name_or_path BAAI/bge-code-v1 \
+    --embedder_model_class decoder-only-base \
+    --query_instruction_format_for_retrieval '<instruct>{}\n<query>{}' \
+    --embedder_query_max_length 2048 \
+    --embedder_passage_max_length 2048 \
+    --trust_remote_code True \
+    --pooling_method last_token \
+    --embedder_batch_size 64 \
+    --devices cuda:0 cuda:1 cuda:2 cuda:3 cuda:4 cuda:5 cuda:6 cuda:7 \
+    --tasks apps codetrans-contest codetrans-dl cosqa synthetic-text2sql stackoverflow-qa codefeedback-mt codefeedback-st CodeSearchNet-ccr-go CodeSearchNet-ccr-java CodeSearchNet-ccr-javascript CodeSearchNet-ccr-php CodeSearchNet-ccr-python CodeSearchNet-ccr-ruby CodeSearchNet-go CodeSearchNet-java CodeSearchNet-javascript CodeSearchNet-php CodeSearchNet-python CodeSearchNet-ruby \
+    --cache_dir ./cache

research/BGE_Coder/evaluation/coir_eval/main.py

@@ -0,0 +1,188 @@
+import os
+import json
+import coir
+from transformers import HfArgumentParser
+
+from arguments import COIREvalArgs, COIREvalModelArgs
+from prompts import get_task_def_by_task_name
+from FlagEmbedding import FlagLLMModel, FlagModel, FlagPseudoMoEModel
+
+
+def get_model(model_args: COIREvalModelArgs):
+    embedder_name_or_path = model_args.embedder_name_or_path
+
+    if model_args.embedder_model_class == "encoder-only-base":
+        embedder = FlagModel(
+            model_name_or_path=embedder_name_or_path,
+            normalize_embeddings=model_args.normalize_embeddings,
+            pooling_method=model_args.pooling_method,
+            use_fp16=model_args.use_fp16,
+            use_bf16=model_args.use_bf16,
+            query_instruction_for_retrieval=model_args.query_instruction_for_retrieval,
+            query_instruction_format=model_args.query_instruction_format_for_retrieval,
+            devices=model_args.devices,
+            trust_remote_code=model_args.trust_remote_code,
+            cache_dir=model_args.cache_dir,
+            batch_size=model_args.embedder_batch_size,
+            query_max_length=model_args.embedder_query_max_length,
+            passage_max_length=model_args.embedder_passage_max_length,
+        )
+    elif model_args.embedder_model_class == "decoder-only-base":
+        embedder = FlagLLMModel(
+            model_name_or_path=embedder_name_or_path,
+            normalize_embeddings=model_args.normalize_embeddings,
+            pooling_method=model_args.pooling_method,
+            use_fp16=model_args.use_fp16,
+            use_bf16=model_args.use_bf16,
+            query_instruction_for_retrieval=model_args.query_instruction_for_retrieval,
+            query_instruction_format=model_args.query_instruction_format_for_retrieval,
+            devices=model_args.devices,
+            examples_for_task=model_args.examples_for_task,
+            examples_instruction_format=model_args.examples_instruction_format,
+            trust_remote_code=model_args.trust_remote_code,
+            cache_dir=model_args.cache_dir,
+            batch_size=model_args.embedder_batch_size,
+            query_max_length=model_args.embedder_query_max_length,
+            passage_max_length=model_args.embedder_passage_max_length,
+        )
+    elif model_args.embedder_model_class == "decoder-only-pseudo_moe":
+        embedder = FlagPseudoMoEModel(
+            model_name_or_path=embedder_name_or_path,
+            normalize_embeddings=model_args.normalize_embeddings,
+            pooling_method=model_args.pooling_method,
+            use_fp16=model_args.use_fp16,
+            use_bf16=model_args.use_bf16,
+            query_instruction_for_retrieval=model_args.query_instruction_for_retrieval,
+            query_instruction_format=model_args.query_instruction_format_for_retrieval,
+            devices=model_args.devices,
+            examples_for_task=model_args.examples_for_task,
+            examples_instruction_format=model_args.examples_instruction_format,
+            trust_remote_code=model_args.trust_remote_code,
+            cache_dir=model_args.cache_dir,
+            batch_size=model_args.embedder_batch_size,
+            query_max_length=model_args.embedder_query_max_length,
+            passage_max_length=model_args.embedder_passage_max_length,
+            domain_for_pseudo_moe=model_args.domain_for_pseudo_moe,
+        )
+    else:
+        raise ValueError(f"Invalid model class: {model_args.embedder_model_class}")
+    embedder.model.config._name_or_path = model_args.embedder_name_or_path
+
+    class CustomFlagModel:
+        def __init__(self, model):
+            self.model = model
+
+        def encode_queries(self, queries, show_progress_bar, convert_to_tensor, **kwargs):
+            if isinstance(queries, str):
+                queries = [queries]
+
+            if isinstance(queries[0], dict):
+                queries = [(e.get('title') + ' ' + e['text']).strip() for e in queries]
+
+            return self.model.encode_queries(queries, **kwargs)
+
+        def encode_corpus(self, corpus, show_progress_bar, convert_to_tensor, **kwargs):
+            if isinstance(corpus, str):
+                corpus = [corpus]
+
+            if isinstance(corpus[0], dict):
+                corpus = [(e.get('title') + ' ' + e['text']).strip() for e in corpus]
+
+            return self.model.encode_corpus(corpus, **kwargs)
+
+        def encode(self, corpus, show_progress_bar, convert_to_tensor, **kwargs):
+            if isinstance(corpus, str):
+                corpus = [corpus]
+
+            if isinstance(corpus[0], dict):
+                corpus = [(e.get('title') + ' ' + e['text']).strip() for e in corpus]
+
+            return self.model.encode(corpus, **kwargs)
+
+    return CustomFlagModel(embedder)
+
+
+def main(
+    eval_args: COIREvalArgs,
+    model_args: COIREvalModelArgs
+):
+    model = get_model(model_args)
+
+    output_folder = os.path.join(eval_args.output_dir, os.path.basename(model.model.model.config._name_or_path))
+
+    all_task = eval_args.tasks
+    if not isinstance(all_task, list):
+        all_task = [all_task]
+
+    all_results = {}
+    for task_name in all_task:
+        save_path = os.path.join(output_folder, f"{task_name}.json")
+        if os.path.exists(save_path):
+            with open(save_path, "r", encoding="utf-8") as f:
+                results = json.load(f)
+                all_results[task_name] = results['metrics']
+                continue
+
+        tmp_task = coir.get_tasks(tasks=[task_name])
+        evaluation = coir.COIR(tasks=tmp_task,
+                               batch_size=model_args.embedder_batch_size)
+
+        model.model.stop_self_pool()
+
+        if eval_args.use_special_instructions:
+            model.model.query_instruction_for_retrieval = get_task_def_by_task_name(task_name)
+
+        results = evaluation.run(model, output_folder=output_folder)
+        all_results[task_name] = results[task_name]
+
+    csn_result = 0
+    csn_num = 0
+    csn_ccr_result = 0
+    csn_ccr_num = 0
+    pop_keys = []
+    all_result = 0
+    all_num = 0
+    for k in all_results.keys():
+        if 'CodeSearchNet-ccr' in k:
+            csn_ccr_result += all_results[k]['NDCG']['NDCG@10']
+            csn_ccr_num += 1
+            pop_keys.append(k)
+        elif 'CodeSearchNet' in k:
+            csn_result += all_results[k]['NDCG']['NDCG@10']
+            csn_num += 1
+            pop_keys.append(k)
+        else:
+            all_result += all_results[k]['NDCG']['NDCG@10']
+            all_num += 1
+    if csn_num > 0:
+        print('Using CodeSearchNet')
+        all_result += csn_result / csn_num
+        all_num += 1
+    if csn_ccr_num > 0:
+        print('Using CodeSearchNet-ccr')
+        all_result += csn_ccr_result / csn_ccr_num
+        all_num += 1
+    new_results = {}
+    for k in all_results:
+        if k in pop_keys:
+            continue
+        new_results[k] = all_results[k]['NDCG']['NDCG@10']
+    if csn_num > 0:
+        new_results['CodeSearchNet'] = csn_result / csn_num
+    if csn_ccr_num > 0:
+        new_results['CodeSearchNet_CCR'] = csn_ccr_result / csn_ccr_num
+    new_results['all'] = all_result / all_num
+
+    print(new_results)
+
+    with open(os.path.join(output_folder, 'OVERALL-results.json'), 'w') as f:
+        json.dump(new_results, f)
+
+
+if __name__ == "__main__":
+    parser = HfArgumentParser((
+        COIREvalArgs,
+        COIREvalModelArgs
+    ))
+    eval_args, model_args = parser.parse_args_into_dataclasses()
+    main(eval_args, model_args)

research/BGE_Coder/evaluation/coir_eval/prompts.py

@@ -0,0 +1,38 @@
+from typing import Dict
+
+
+def get_task_def_by_task_name(task_name: str) -> str:
+    task_name_to_instruct: Dict[str, str] = {
+        # Text-to-Code Retrieval
+        ## Code Contest Retrieval
+        'apps': 'Given a code contest problem description, retrieve relevant code that can help solve the problem.',
+        ## Web Query to Code Retrieval
+        'cosqa': 'Given a web search query, retrieve relevant code that can help answer the query.',
+        ## Text-to-SQL Retrieval
+        'synthetic-text2sql': 'Given a question in text, retrieve SQL queries that are appropriate responses to the question.',
+        
+        # Code-to-Text Retrieval
+        ## Code Summary Retrieval
+        'CodeSearchNet-': 'Given a piece of code, retrieve the document string that summarizes the code.',
+        
+        # Code-to-Code Retrieval
+        ## Code Context Retrieval
+        'CodeSearchNet-ccr-': 'Given a piece of code segment, retrieve the code segment that is the latter part of the code.',
+        ## Similar Code Retrieval
+        'codetrans-dl': 'Given a piece of code, retrieve code that is semantically equivalent to the input code.',
+        'codetrans-contest': 'Given a piece of Python code, retrieve C++ code that is semantically equivalent to the input code.',
+        
+        # Hybrid Code Retrieval
+        ## Single-turn Code QA
+        'stackoverflow-qa': 'Given a question that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.',
+        'codefeedback-st': 'Given a question that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.',
+        ## Multi-turn Code QA
+        'codefeedback-mt': 'Given a multi-turn conversation history that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.',
+    }
+    
+    special_task_names = ['CodeSearchNet-ccr-', 'CodeSearchNet-']
+    for special_task_name in special_task_names:
+        if special_task_name in task_name:
+            return task_name_to_instruct[special_task_name]
+    
+    return task_name_to_instruct[task_name]