import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline


title = "Code Compexity Predictor [WIP]"
description = "This is a space to predict complexity of Java code with [CodeParrot-Multi-Complexity](https://huggingface.co/codeparrot/codeparrot-small-multi),\
    a multilingual model for code generation, fine-tuned on [CodeComplex](https://huggingface.co/datasets/codeparrot/codecomplex), a dataset for complexity prediction of Java code."

#add examples
example = [
    ["code example 1", "nlogn"],
    ["code example 2", "constant"]]

# model to be changed to the finetuned one
tokenizer = AutoTokenizer.from_pretrained("codeparrot/codeparrot-small-multi")
model = AutoModelForSequenceClassification.from_pretrained("codeparrot/codeparrot-small-multi", num_labels=7)


def complexity_estimation(gen_prompt, topk):
    pipe = pipeline("text-classification", model=model, tokenizer=tokenizer)
    output = pipe(gen_prompt)[0]
    # add label conversion to class
    label = output['label']
    score = output['score']
    return label, score


iface = gr.Interface(
    fn=complexity_estimation, 
    inputs=[
        gr.Textbox(lines=10, label="Input code"),
        gr.inputs.Slider(
            minimum=1,
            maximum=3,
            step=1,
            default=1,
            label="Number of results to return",
        ),
    ],
    outputs=[
    gr.Textbox(label="Predicted complexity", lines=1) ,
    gr.Textbox(label="Corresponding probability", lines=1) ,
],
    examples=example,
    layout="vertical",
    theme="peach",
    description=description,
    title=title
)
iface.launch()