refactor: move functions from utils.data to data

codecritic/cli/mk_cov_dataset.py

@@ -7,7 +7,7 @@ from codecritic.data.cov import (
 )
 
 from codecritic.utils.json import load_json
-from codecritic.utils.data import save_jsonl_dataset
+from codecritic.data.utils import save_jsonl_dataset
 from codecritic.utils.vllm import vllm_chatcomplete
 
 

codecritic/cli/mk_rm_dataset.py

 import argparse
 from pathlib import Path
 from codecritic.utils.json import load_json
-from codecritic.utils.data import save_jsonl_dataset
+from codecritic.data.utils import save_jsonl_dataset
 
 from codecritic.data.edit_distance import (
     mk_problem_groups,

codecritic/cli/run_sft_test.py

@@ -2,12 +2,15 @@ import argparse
 from pathlib import Path
 import pprint
 
+from codecritic.data.utils import mk_message
+from codecritic.data.verify import JUDGE_PROMPT
 from transformers import AutoTokenizer
 
+from codecritic.data.code import extract_code, code_template
 from codecritic.data.cov import COV_PROMPT
 from codecritic.utils.vllm import vllm_chatcomplete, vllm_score
 from codecritic.utils.json import load_jsonl, save_jsonl 
-from codecritic.utils.data import extract_code, code_template, mk_message, JUDGE_PROMPT, get_score_token_id
+from codecritic.data.verify import get_score_token_id
 from codecritic.utils.metric import group_results, score_pass_at_k
 
 

codecritic/cli/verify_rm_to_sft.py

@@ -4,8 +4,9 @@
 # 2. Using SFT (Supervised Fine-Tuning) directly
 # This experiment aims to fairly compare these two approaches.
 import argparse
+from codecritic.data.utils import mk_message, mk_messages, save_jsonl_dataset
 from codecritic.utils.json import load_json
-from codecritic.utils.data import mk_message, mk_critic_verify, mk_messages, save_jsonl_dataset
+from codecritic.data.verify import mk_critic_verify
 
 
 def convert_preference_to_sft(item):

codecritic/data/cov.py

@@ -2,11 +2,9 @@
 # Is reasoning really work? Let's verify step by step.
 from codecritic.data.code import extract_code, code_template
 
-from codecritic.utils.data import (
-    mk_message,
-    mk_messages,
+from codecritic.data.utils import SPLITTER, mk_message, mk_messages
+from codecritic.data.verify import (
     mk_critic_verify,
-    SPLITTER,
 )
 
 COV_PROMPT = "Please verify your code step by step using Markdown code blocks. After each step, explain whether it's correct or not, and if not, explain the issue."

codecritic/utils/data.py → codecritic/data/utils.py

-from pathlib import Path
-
+# Note that the human and observation should appear in odd positions
+# while llm should appear in even positions.
 from codecritic.utils.json import save_jsonl
 
 
-# Note that the human and observation should appear in odd positions
-# while llm should appear in even positions.
+from pathlib import Path
+
+
 def mk_messages(messages):
     return {"messages": messages}
 
@@ -16,29 +17,10 @@ def mk_message(user, assistant):
     ]
 
 
-JUDGE_PROMPT = "Is the code correct (Yes/No)?"
-
-
-def mk_critic_verify(answer=None):
-    # answer: bool or none
-    message = [{"role": "user", "content": JUDGE_PROMPT}]
-    if answer is not None:
-        response = "Yes" if answer else "No"
-        message.append({"role": "assistant", "content": response})
-
-    return message
-
-
 def save_jsonl_dataset(dataset, output_dir, split="train"):
     output_dir = Path(output_dir)
     output_dir.mkdir(parents=True, exist_ok=True)
     save_jsonl(dataset, output_dir / f"{split}.jsonl")
 
 
-def get_score_token_id(tokenizer, token_str="Yes"):
-    score_tokens = tokenizer.encode(token_str, add_special_tokens=False)
-    assert len(score_tokens) == 1
-    return score_tokens[0]
-
-
 SPLITTER = "__I_wish_it_were_weekends_all_the_time.__"

codecritic/data/verify.py

+JUDGE_PROMPT = "Is the code correct (Yes/No)?"
+
+
+def mk_critic_verify(answer=None):
+    # answer: bool or none
+    message = [{"role": "user", "content": JUDGE_PROMPT}]
+    if answer is not None:
+        response = "Yes" if answer else "No"
+        message.append({"role": "assistant", "content": response})
+
+    return message
+
+
+def get_score_token_id(tokenizer, token_str="Yes"):
+    score_tokens = tokenizer.encode(token_str, add_special_tokens=False)
+    assert len(score_tokens) == 1
+    return score_tokens[0]

codecritic/utils/vllm.py

@@ -5,7 +5,7 @@ import multiprocessing
 from itertools import chain
 from functools import partial
 
-from codecritic.utils.data import SPLITTER
+from codecritic.data.utils import SPLITTER
 import numpy as np
 
 def generate_worker(cuda_device, prompts, model_path, sampling_params):