이것은 "CodeRAG-Bench: Retrieval이 코드 생성(Code Generation)을 증강할 수 있는가?" 프로젝트를 위한 코드 저장소(repository)입니다.

새로운 환경을 생성하세요:

conda env create -n crag python=3.10 -y
conda activate crag

그리고 필요한 라이브러리들을 설치하세요:

pip install -r requirements.txt

• Retrieval (검색): BM25, sentence-transformers를 통한 밀집 검색기 (dense retrievers), 그리고 독점 API 임베딩 (proprietary API embeddings)을 사용하여 검색을 실행하는 코드입니다.
• Generation (생성): 모델 생성 및 실행 기반 평가 (execution-based evaluation)를 실행하는 코드입니다.
• Preprocess (전처리): 검색 풀 (retrieval pool) 구축을 위해 원시 데이터를 전처리하는 코드이며, 자세한 내용은 디렉토리 내부를 참조하세요.

cd retrieval/

데이터셋에 대해 검색을 실행하기 전에, 해당 데이터셋을 위한 데이터스토어 (datastore)를 생성해야 합니다. 다음을 따르세요:

python -m create/${data_name}.py
# ${data_name} 선택지
# 기본 프로그래밍: 'humaneval', 'mbpp', 'live_code_bench'
...

만약 ModuleNotFoundError가 발생한다면,

PYTHONPATH=./ python create/${data_name}.py 명령어로 실행해 보세요.

새로운 데이터셋을 실행하려면, 데이터를 BEIR 공식 형식으로 재구성하기만 하면 됩니다. 이 형식은 corpus.jsonl, queries.jsonl, 그리고 qrel/test.tsv라는 다음 세 가지 파일을 포함하는 데이터셋 디렉토리를 생성합니다.

Code Search Net을 문서-대-코드 검색 (document-to-code retrieval) 작업으로 변환하는 예시 스크립트는 code_search_net.py를 참조하세요.
결과물인 데이터셋이 retrieval/datasets 아래에 있는지 확인하세요.

다음과 같이 sentence-transformers에서 임베딩 모델을 로드하여 임베딩 모델을 실행하세요:

python3 eval_beir_sbert_canonical.py \
--model YOUR_MODEL_NAME_OR_PATH \
--dataset TASK_NAME \
...

출력 파일 이름인 --output_file을 지정함으로써, 검색 결과를 json 파일로 저장할 수 있습니다.

매개변수 --output_file을 지정함으로써, 검색 결과를 json 파일로 저장할 수 있습니다.

{'ndcg': {'NDCG@1': 0.61667, 'NDCG@3': 0.68203, 'NDCG@5': 0.70804, 'NDCG@10': 0.72701, 'NDCG@100': 0.74926, 'NDCG@1000': 0.75551}, 'mrr': {'MRR@1': 0.61667, 'MRR@3': 0.67278, 'MRR@5': 0.68611, 'MRR@10': 0.69368, 'MRR@100': 0.69721, 'MRR@1000': 0.69744}, 'recall': {'Recall@1': 0.58817, 'Recall@3': 0.728, 'Recall@5': 0.79294, 'Recall@10': 0.84789, 'Recall@100': 0.95, 'Recall@1000': 0.99667}, 'precision': {'P@1': 0.61667, 'P@3': 0.26444, 'P@5': 0.176, 'P@10': 0.09533, 'P@100': 0.01077, 'P@1000': 0.00113}}```

`--results_file`은 검색 결과를 저장할 파일 이름을 지정하며, 이는 이후의 RAG 평가에 사용됩니다.

오픈 검색(open retrieval)의 경우, 저희 huggingface 스페이스에서 코퍼스 파일(list of the corpora)을 로드하고 단일 또는 다중 GPU를 사용하여 임베딩을 생성할 수 있습니다.

- 단일 GPU를 사용하여 임베딩 생성

python generate_embeddings.py
--model YOUR_MODEL_NAME_OR_PATH
--output_dir OUTPUT_EMBEDDING_DIR
...

- 다중 GPU를 사용하여 임베딩 생성 (예: 8)

for i in {0..7}; do
export CUDA_VISIBLE_DEVICES=${i}
nohup python generate_embeddings.py --model_name_or_path YOUR_MODEL_NAME_OR_PATH
...

이제 생성된 임베딩을 로드하여 대상 데이터셋에 대한 오픈 검색을 실행할 수 있습니다.

python3 eval_beir_sbert_open.py
--model avsolatorio/GIST-large-Embedding-v0
--embdding_path

python3 modify_corpus_for_bm25.py \
--dataset DATASET_NAME, 모든 데이터셋에 대해 작업을 수행하려면 "all" 입력 \
--output_metadir OUTPUT_DIR \
...

OUTPUT_DIR/{DATASET_NAME}_corpus/edit.jsonl로부터 코퍼스 (corpus) 인덱싱 (Indexing)을 수행하며,

인덱스는 INDEX_DIR/{DATASET_NAME}_corpus/에 저장됩니다.

:

python3 modify_corpus_for_bm25.py \
--dataset DATASET_NAME, 모든 데이터셋에 대해 작업을 수행하려면 "all" 입력 \
--output_metadir OUTPUT_DIR \
...

BM25를 사용하여 대상 데이터셋에서 쿼리 (query)를 검색합니다:

python3 modify_corpus_for_bm25.py \
--dataset DATASET_NAME, 모든 데이터셋에 대해 작업을 수행하려면 "all" 입력 \
--output_metadir OUTPUT_DIR \
...

점수 파일은 results/{DATASET_NAME}_k1={K1}_b={B}_pyserini_bm25_output.jsonl에 저장되며,

검색 결과는 results/{DATASET_NAME}_k1={K1}_b={B}_pyserini_bm25.jsonl에 저장됩니다.

사용자의 편의를 위해 --stage all과 그에 해당하는 파라미터 (parameters)를 전달하여 모든 단계를 한 번에 실행할 수 있습니다.

각 인스턴스 (instance)의 코퍼스를 수정, 인덱싱 및 검색합니다.

python eval_beir_pyserini_repo.py \
--dataset DATASET_NAME \
--output_metadir OUTPUT_DIR \
...

수정된 코퍼스는 OUTPUT_DIR/{DATASET_NAME}_corpus/{instance dir}/edit.jsonl에 저장되며,

인덱스는 INDEX_DIR/{DATASET_NAME}_corpus/{instance dir}에 저장됩니다.

이 메타 스크립트 (meta script)에서는 swe-bench-lite 이외의 데이터셋도 지원합니다:

python3 eval_corpora_ablations.py \
--model bm25 \
--dataset DATASET_NAME, 모든 데이터셋에 대해 작업을 수행하려면 "all" 입력 \
...

수정된 코퍼스는 OUTPUT_DIR/{CORPUS NAME}_corpus/edit.jsonl에 저장되며,

인덱스는 INDEX_DIR/{CORPUS NAME}_corpus/에 저장됩니다. 검색 결과는 results/{DATASET_NAME}_k1={K1}_b={B}_pyserini_bm25_corpus에 저장되며,

점수 파일은 {DATASET_NAME}_corpus={CORPUS_NAME}_k1={K1}_b={B}_pyserini_bm25_output.jsonl에, 검색 결과는 {DATASET_NAME}_corpus={CORPUS_NAME}_k1={K1}_b={B}_pyserini_bm25.jsonl에 저장됩니다.

리포지토리 레벨(repository level)이 아닌 데이터셋의 경우: 다음과 같이 독점 API(proprietary APIs)로부터 임베딩 (embeddings)을 로드하여 API 기반 모델을 실행하십시오:

voyage.ai

python3 eval_voyage.py \ --dataset TASK_NAME \ --model MODEL_NAME (기본값은 voyage-code-2) \ --api_key_fp PATH_TO_YOUR_API_KEY_FILE (새 줄이 필요함) \ --batch_size YOUR_BATCH_SIZE \ --output_file PATH_TO_YOUR_SCORE_FILE \ --results_file PATH_TO_YOUR_RETRIEVAL_RESULTS_FILE

openai

python3 eval_openai.py \ --dataset TASK_NAME \ --model MODEL_NAME (기본값은 text-embedding-3-small) \ --api_key_fp PATH_TO_YOUR_API_KEY_FILE (새 줄이 필요함) \ --batch_size YOUR_BATCH_SIZE \ --output_file PATH_TO_YOUR_SCORE_FILE \ --results_file PATH_TO_YOUR_RETRIEVAL_RESULTS_FILE

--run_async 옵션은 검색 (retrieval)을 비동기적으로 실행하는 데 사용할 수 있습니다. 기본 동작은 voyage의 경우 생성된 문서 임베딩 (document embedding) 및 문서 ID (doc ids)를 다음 위치에 캐싱하여 사용합니다:

datasets/{dataset}/voyage_doc_embeddings.npy

및

datasets/{dataset}/voyage_doc_ids.json

openai의 경우 다음 위치에 캐싱합니다:

datasets/{dataset}/doc_embeddings.npy

및

datasets/{dataset}/doc_ids.json

쿼리 임베딩 (Query embeddings) 및 쿼리-ID 매핑 (query to ids) 또한 다음 위치에 캐싱됩니다:

datasets/{dataset}/voyage_query_embeddings.npy

및

datasets/{dataset}/voyage_queryidx2truncatedidx.json

openai의 경우 다음 위치에 캐싱됩니다:

datasets/{dataset}/query_embeddings.npy

및

datasets/{dataset}/queryidx2truncatedidx.json

새롭게 시작하고 싶다면 이 파일들을 삭제하십시오.

리포지토리 레벨 (repository level) 데이터셋의 경우:

python3 eval_api_repo.py \ --dataset TASK_NAME \ --model MODEL_NAME \ --api_key_fp PATH_TO_YOUR_API_KEY_FILE (새 줄이 필요함) \ --batch_size YOUR_BATCH_SIZE \ --output_file PATH_TO_YOUR_SCORE_FILE \ --results_file PATH_TO_YOUR_RETRIEVAL_RESULTS_FILE

--run_async 옵션은 검색 (retrieval)을 비동기적으로 실행하는 데 사용할 수 있습니다. 쿼리 임베딩 (query embeddings), 쿼리-ID 매핑 (query to id) 및 문서 임베딩 (document embeddings)과 문서 ID (document ids)는 모두 datasets/{dataset}/{instance dir}에 저장되며, 여기서 instance dir은 데이터셋의 특정 인스턴스 (instance)를 의미합니다.

swe-bench-lite 이외의 데이터셋도 지원됩니다.
위의 명령어들에 --corpus_path THE_PATH_TO_YOUR_CORPUS_FILE (별도의 디렉토리에 있는 edit.jsonl 파일을 기대하며, BM25를 위해 처리된 파일을 사용할 수 있습니다))를 추가하세요.

쿼리 임베딩 (Query embeddings) 및 쿼리 ID (query to ids)는 위와 동일한 방식으로 캐싱(cached)되는 반면, 문서는 다음과 같이 캐싱됩니다:

Voyage의 경우 THE_PARENT_DIR_OF_CORPUS_PATH/voyage_doc_embeddings.npy 및 THE_PARENT_DIR_OF_CORPUS_PATH/voyage_doc_ids.json에 생성된 문서 임베딩 (document embedding)과 doc ID가 저장됩니다.

OpenAI의 경우 THE_PARENT_DIR_OF_CORPUS_PATH/doc_embeddings.npy 및 THE_PARENT_DIR_OF_CORPUS_PATH/doc_ids.json에 저장됩니다.

새롭게 시작하고 싶다면 이 파일들을 삭제해 주세요.

main.py 스크립트는 Hugging Face 또는 OpenAI에서 지원하는 모든 모델을 사용하여 코드 생성 (code generation)을 실행할 수 있습니다.

원본 데이터셋에 대해 검색 없는 생성 (no-retrieval generation)을 실행하려면, dataset_path 인자에 해당 Hugging Face 데이터셋 이름을 지정하세요:

python main.py --task "humaneval" \
--model "bigcode/starcoder2-7b" \
--dataset_path "openai_humaneval" \
...

코드 실행 (code execution)을 통해 생성된 결과를 평가하려면 --allow_code_execution을 설정하세요. 이는 모든 태스크 (tasks)에 필수적입니다.

task는 dataset_path와 일치해야 함에 유의하세요. 사용 가능한 모든 태스크는 다음과 같습니다:

• 기초 프로그래밍 (basic programming): 'humaneval', 'mbpp', 'lcb' (LiveCodeBench용)
• 오픈 도메인 (open domain): 'ds1000-all-completion', 'odex-en'
• 리포지토리 레벨 (repository level): 'repoeval-function', 'swebench-lite'

이전 검색 결과(예: "retrieval/humaneval/gist_large.json")를 사용하여 생성을 실행하려면, 다음과 같이 파일을 지정하세요:

python main.py --task "humaneval" \
--model "bigcode/starcoder2-7b" \
--dataset_path "json" --data_files_test "retrieval/humaneval/gist_large.json" \
...

main.py 스크립트를 실행하면 생성이 완료된 후 자동으로 실행 기반 평가 (execution-based evaluation)를 수행합니다.
하지만 RepoEval(-function) 및 SWE-bench(-Lite) 데이터셋의 경우, 문제의 복잡성으로 인해 추가적인 설정이 필요합니다.

RepoEval을 위한 저장소(repositories)를 다운로드하고 (예: retrieval/output/repoeval/repositories/function_level/ 아래), 검색 결과(retrieval results)를 얻은 후 (예: /path/to/retrieval/results/retriever-name.jsonl), 이전 단계에서 코드 생성 출력물(code generation outputs)을 확보했다면 (예: /path/to/generation/outputs/model-name.json), 실행 기반 평가(execution-based evaluation)를 수행할 수 있습니다.

먼저, repoeval이라는 환경을 구축합니다. 실험을 실행하기 위해 별도의 환경을 사용할 수 있음에 유의하세요. 실행 시 자동으로 repoeval 환경으로 전환됩니다.

cd generation
conda env create --file eval/tasks/custom_metrics/repoeval_environment.yml -n repoeval

그 다음, 환경을 테스트하기 위해 다음 명령어를 실행합니다. 이 명령어는 모든 저장소에서 테스트를 실행하여 원본 코드(즉, 정답 코드(ground truth code))가 모든 테스트 케이스를 통과할 수 있는지 확인합니다.

cd generation
PYTHONPATH=./ python eval/tasks/custom_metrics/repoeval_execution.py

모든 테스트를 통과하면, Pass@1을 사용하여 코드 생성 출력물을 평가할 수 있습니다.

cd generation
MODEL_NAME="model-name"
RETRIEVAL_FILE="/path/to/retrieval/results/retriever-name.jsonl"
...

이전 단계에서 생성된 출력 파일(예: /path/to/generation/outputs/model-name.json)을 확보한 후, SWE-Bench 평가 하네스(evaluation harness)를 위해 해당 파일을 제자리에서(in-place) 변환하려면 다음을 실행하세요:

python generation/eval/tasks/custom_metrics/swebench_transform.py \
--output_path /path/to/generation/outputs/model-name.json

그 다음, OpenDevin에서 제공하는 SWE-bench 평가용 Docker를 실행합니다 (주의: 30GB 이상의 매우 큰 크기입니다).

docker run -it \
-v /path/to/generation/outputs:/swe_bench_output \
ghcr.io/opendevin/eval-swe-bench:full-v1.0 /bin/bash

그리고 Docker 내부에서 평가를 실행하여 최종 평가 통계(evaluation stats)를 얻습니다:

export MINICONDA3=/swe_util/miniforge3
export OD_SWE_BENCH=/swe_util/OD-SWE-bench
export EVAL_DATA_DIR=/swe_util/eval_data
...

저희는 또한 더 나은 RACG를 위해 문서 청킹 (Document Chunking) 및 재순위화 (Reranking)를 탐색했습니다.

상위 100개의 검색된 문서들을 재순위화하여 더 나은 top-k (k<<100) 문서를 얻기 위해, 기존의 재순위화 (Reranker) 모델을 사용합니다:

python rerank.py --results_path ${retrieval_results_file}

이 스크립트는 results_path에서 사용 가능한 옵션 목록을 제공하며 쿼리 필드 (Query field)를 묻는 프롬프트를 표시합니다.

예를 들어, humaneval 검색 결과의 경우 다음과 같이 프롬프트가 나타납니다:

Choose query field from [dict_keys(['task_id', 'prompt', 'canonical_solution', 'test', 'entry_point', 'docs'])]:

"prompt"를 입력하고 Enter를 누르면 재순위화 (Reranking)가 자동으로 완료됩니다.

다른 데이터셋의 경우, mbpp 및 livecodebench에는 "text"를, ds-1000에는 "prompt"를, odex에는 "intent"를, repoeval에는 "prompt"를, 그리고 swebench에는 "problem_statement"를 사용하십시오.

검색된 문서에서 처음 N-토큰 청크 (N-token chunk)를 가져와 새로운 검색 파일로 처리하려면 다음을 실행하십시오:

python chunk.py --results_path ${retrieval_results_file} --max_num_tokens 500