codeparrot/code-complexity-predictor

1import gradio as gr2from datasets import ClassLabel3from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline4 5 6title = "BigO"7description = "In this space we predict the complexity of Java code with [UniXcoder-java-complexity-prediction](https://huggingface.co/codeparrot/unixcoder-java-complexity-prediction),\8    a multilingual model for code, fine-tuned on [CodeComplex](https://huggingface.co/datasets/codeparrot/codecomplex), a dataset for complexity prediction of Java code."9 10#add examples11example = [['int n = 1000;\nSystem.out.println("Hey - your input is: " + n);'],12    ['class GFG {\n \n    public static void main(String[] args)\n    {\n        int i, n = 8;\n        for (i = 1; i <= n; i++) {\n            System.out.printf("Hello World !!!\n");\n        }\n    }\n}'],13    ['import java.io.*;\nimport java.util.*;\n\npublic class C125 {\n\tpublic static void main(String[] args) throws IOException {\n\t\tBufferedReader r = new BufferedReader(new InputStreamReader(System.in));\n\t\tString s = r.readLine();\n\t\tint n = new Integer(s);\n\t\tSystem.out.println("0 0 "+n);\n\t}\n}\n']]14 15# model to be changed to the finetuned one16tokenizer = AutoTokenizer.from_pretrained("codeparrot/unixcoder-java-complexity-prediction")17model = AutoModelForSequenceClassification.from_pretrained("codeparrot/unixcoder-java-complexity-prediction", num_labels=7)18 19def get_label(output):20    label = int(output[-1])21    labels = ClassLabel(num_classes=7, names=['constant', 'cubic', 'linear', 'logn', 'nlogn', 'np', 'quadratic'])22    return labels.int2str(label)23    24def complexity_estimation(gen_prompt):25    pipe = pipeline("text-classification", model=model, tokenizer=tokenizer)26    output = pipe(gen_prompt)[0]27    # add label conversion to class28    label = get_label(output['label'])29    score = output['score']30    return label, score31 32 33iface = gr.Interface(34    fn=complexity_estimation, 35    inputs=[36        gr.Code(lines=10, language="python", label="Input code"),37    ],38    outputs=[39    gr.Textbox(label="Predicted complexity", lines=1) ,40    gr.Textbox(label="Corresponding probability", lines=1) ,41],42    examples=example,43    layout="vertical",44    theme="darkpeach",45    description=description,46    title=title47)48iface.launch()