Update app.py

app.py

@@ -11,13 +11,16 @@ example = [
 
 # model to be changed to the finetuned one
 tokenizer = AutoTokenizer.from_pretrained("codeparrot/codeparrot-small-multi")
-model = AutoModelForSequenceClassification.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)
-    label = pipe(gen_prompt, do_sample=True, top_k=topk)[0]['label']
-    return label
+    output = pipe(gen_prompt)[0]
+    # add label conversion to class
+    label = output['label']
+    score = output['score']
+    return label, score
 
 
 iface = gr.Interface(
@@ -32,7 +35,10 @@ iface = gr.Interface(
             label="Number of results to return",
         ),
     ],
-    outputs=gr.Textbox(label="Predicted complexity", lines=10),
+    outputs=[
+    gr.Textbox(label="Predicted complexity", lines=1) ,
+    gr.Textbox(label="Corresponding probability", lines=1) ,
+],
     examples=example,
     layout="vertical",
     theme="peach",