luminlemon/codescorer

BTScorer - Performance Dimension Code Scorer

BTScorer scores code along four performance optimization dimensions: CPU, IO, Memory, and Time.

The model outputs a 4-dimensional score vector. A higher score in a dimension means the input code is predicted to be more optimized for that dimension than another input with a lower score.

BTScorer was trained on patch-style inputs. The recommended inference workflow is therefore:

1. Convert a diff or code snippet into the neutral patch-style format.
2. Score the resulting after_code.

1. Environment Setup

Python 3.11 is recommended.

conda create -n codescorer python=3.11
conda activate codescorer
pip install -r requirements.txt

Minimum packages:

pip install torch transformers numpy huggingface_hub

2. Diff JSONL Input

Use this when your data contains unified diffs, for example records like performance_records_with_patches.jsonl:

{
  "id": "example-1",
  "patches": [
    {
      "diff": "diff --git a/input.py b/input.py\n@@ -1,5 +1,4 @@ def total(xs):\n..."
    }
  ]
}

Prepare scoring inputs:

python etl/prepare_scoring_inputs.py \
  --mode diff \
  --input_file performance_records_with_patches.jsonl \
  --output_file scoring_inputs.jsonl \
  --keep_fields title url

The ETL keeps only the after side of each diff, because this scorer scores the code state being evaluated. It converts each hunk into:

# <FILE> path/from/diff.py
# <HUNK> @@ def function_or_class_context(...): @@
<after-code lines from the hunk>

If the original diff hunk has no function/class context, the hunk marker is:

# <HUNK> @@ @@

Line numbers and hunk lengths such as -80,8 +82,7 are removed to avoid length leakage.

Then score the prepared JSONL:

python model/score_code.py \
  --jsonl_file scoring_inputs.jsonl \
  --jsonl_output scored_outputs.jsonl \
  --text_field after_code \
  --input_format already_wrapped

Each output row contains a scores field:

{
  "id": "example-1",
  "after_code": "...",
  "scores": {
    "CPU": 1.23,
    "IO": 0.04,
    "Memory": 0.87,
    "Time": 1.11
  }
}

3. Plain Code JSONL Input

Use this when your data contains ordinary code snippets rather than diffs:

{"id": "snippet-1", "code": "def total(xs):\n    return sum(xs)\n"}

Prepare scoring inputs:

python etl/prepare_scoring_inputs.py \
  --mode code \
  --input_file code_snippets.jsonl \
  --output_file scoring_inputs.jsonl \
  --code_field code

The ETL wraps each snippet as patch-style input. The default file name is input.py. If the code contains a def, async def, or class line, that line is used as the hunk context:

# <FILE> input.py
# <HUNK> @@ def total(xs): @@
def total(xs):
    return sum(xs)

If no function/class line is found, the wrapper still works:

# <FILE> input.py
# <HUNK> @@ @@
<your code>

If your JSONL has file path or hunk context fields, pass them explicitly:

python etl/prepare_scoring_inputs.py \
  --mode code \
  --input_file code_snippets.jsonl \
  --output_file scoring_inputs.jsonl \
  --code_field code \
  --file_field file_path \
  --hunk_context_field function_name

Then score the prepared JSONL:

python model/score_code.py \
  --jsonl_file scoring_inputs.jsonl \
  --jsonl_output scored_outputs.jsonl \
  --text_field after_code \
  --input_format already_wrapped

4. Single Snippet Usage

For quick tests, you can score one code snippet directly. The script automatically applies the same neutral wrapper.

from model.score_code import BTScorerInference

scorer = BTScorerInference()

code = """
def sum_array(arr):
    total = 0
    for x in arr:
        total += x
    return total
"""

scores = scorer.score(code)
print(scores)

Command line:

python model/score_code.py \
  --code "def total(xs): return sum(xs)"

If you know the file path or hunk context, provide it:

python model/score_code.py \
  --code_file input.py \
  --wrapped_file_name utils/math_ops.py \
  --hunk_context "def total(xs):"

5. Notes

During inference, no label mask is required. The mask was only used during training to calculate the Bradley-Terry loss for specific optimization labels.

The wrapper metadata is additive: it does not remove or rewrite the code body. If a function header appears in the input code, it remains in the code body; the hunk marker simply provides extra patch-style context.

By default, the script uses CUDA when available and loads the base encoder in bfloat16. If CUDA is unavailable, it exits instead of silently loading the 1.5B model on CPU. Use --allow_cpu only when CPU inference is intentional.