import gradio as gr from datasets import ClassLabel from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline title = "BigO" description = "In this space we predict the complexity of Java code with [UniXcoder-java-complexity-prediction](https://huggingface.co/codeparrot/unixcoder-java-complexity-prediction),\ a multilingual model for code, fine-tuned on [CodeComplex](https://huggingface.co/datasets/codeparrot/codecomplex), a dataset for complexity prediction of Java code." #add examples example = [['int n = 1000;\nSystem.out.println("Hey - your input is: " + n);'], ['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}'], ['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']] # model to be changed to the finetuned one tokenizer = AutoTokenizer.from_pretrained("codeparrot/unixcoder-java-complexity-prediction") model = AutoModelForSequenceClassification.from_pretrained("codeparrot/unixcoder-java-complexity-prediction", num_labels=7) def get_label(output): label = int(output[-1]) labels = ClassLabel(num_classes=7, names=['constant', 'cubic', 'linear', 'logn', 'nlogn', 'np', 'quadratic']) return labels.int2str(label) def complexity_estimation(gen_prompt): pipe = pipeline("text-classification", model=model, tokenizer=tokenizer) output = pipe(gen_prompt)[0] # add label conversion to class label = get_label(output['label']) score = output['score'] return label, score iface = gr.Interface( fn=complexity_estimation, inputs=[ gr.Code(lines=10, language="java", label="Input code"), ], outputs=[ gr.Textbox(label="Predicted complexity", lines=1) , gr.Textbox(label="Corresponding probability", lines=1) , ], examples=example, layout="vertical", theme="darkpeach", description=description, title=title ) iface.launch()