app.py

import ast
from pathlib import Path

import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
from torch import nn


model_id = "answerdotai/ModernBERT-base"
path = "DanGalt/modernbert-code-comrel-synthetic"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSequenceClassification.from_pretrained(path)
sep = "[SEP]"


def prepare_input(example):
    tokens = tokenizer(
        example["function_definition"] + sep + example["code"] + sep + example["comment"],
        truncation=True,
        max_length=1024,
        return_tensors="pt"
    )
    return tokens


def parse_text(text):
    # NOTE: Doesn't collect comments and function definitions correctly
    inputs = []
    defs = []
    tree = ast.parse(text)
    for el in tree.body:
        if isinstance(el, ast.FunctionDef):
            defs.append((el.lineno - 1, el.end_lineno - 1, el.col_offset))

    inputs = []
    lines = text.split('\n')
    for lineno, line in enumerate(lines):
        if (offset := line.find('#')) != -1:
            corresponding_def = None
            for (def_l, def_el, def_off) in defs:
                if def_l <= lineno and def_off <= offset:
                    corresponding_def = (def_l, def_el, def_off)

            comment = line[offset:]
            code = '\n'.join(lines[lineno - 4:lineno + 4])
            fdef = "None"
            if corresponding_def is not None:
                fdef = [lines[corresponding_def[0]][corresponding_def[2]:]]
                cur_lineno = corresponding_def[0]
                while cur_lineno <= corresponding_def[1]:
                    if lines[cur_lineno].find("):") != -1 or lines[cur_lineno].find("->") != -1:
                        fdef += lines[corresponding_def[0] + 1:cur_lineno + 1]
                        break
                    cur_lineno += 1

                fdef = '\n'.join(fdef).strip()

            inputs.append({
                "function_definition": fdef,
                "code": code,
                "comment": comment,
                "lineno": lineno
            })
    return inputs


def predict(inp, model=model):
    with torch.no_grad():
        out = model(**inp)
    return nn.functional.softmax(out.logits, dim=-1)[0, 1].item()


def parse_and_predict(text, thrd=0.0):
    parsed = parse_text(text)
    preds = [predict(prepare_input(p)) for p in parsed]
    result = []
    for i, p in enumerate(preds):
        if thrd > 0:
            p = thrd > p
        result.append((parsed[i]["lineno"], p))

    return result


def parse_and_predict_file(path, thrd=0.0):
    text = Path(path).open("r").read()
    return parse_and_predict(text, thrd)


def parse_and_predict_pretty_out(text, thrd=0.0):
    results = parse_and_predict(text, thrd=thrd)
    lines = text.split('\n')
    output = []
    if thrd > 0:
        for lineno, do_warn in results:
            if do_warn:
                output.append(f"The comment on line {lineno} is incorrect: '{lines[lineno]}'.")
    else:
        for lineno, p in results:
            output.append(f"The comment on line {lineno} is estimated to be correct with probability {p:.2f}: '{lines[lineno]}'.")
    return '\n'.join(output)


example_text = """a = 3
b = 2
# The code below does some calculations based on a predefined rule that is very important
c = a - b  # Calculate and store the sum of a and b in c
d = a + b  # Calculate and store the sum of a and b in d
e = c * b  # Calculate and store the product of c and d in e
print(f"Wow, maths: {[a, b, c, d, e]}")"""

gradio_app = gr.Interface(
    fn=parse_and_predict_pretty_out,
    inputs=[
        gr.Textbox(label="Input", lines=7),
        gr.Slider(value=0.8, minimum=0.0, maximum=1.0, step=0.05)],
    outputs=[gr.Textbox(label="Predictions", lines=7)],
    examples=[[example_text, 0.0], [example_text, 0.53]],
    title="Comment \"Correctness\" Classifier",
    description='Calculates probabilities for each comment in text to be "correct"/"relevant". If the threshold is 0, outputs raw predictions. Otherwise, will report only "incorrect" comments.'
)

if __name__ == "__main__":
    gradio_app.launch()