Update app.py

app.py

@@ -8,19 +8,23 @@ description = "This is a space to predict complexity of Java code with [CodeParr
 
 #add examples
 example = [
-    ["code example 1", 1],
-    ["code example 2", 1]]
+    ['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', 1],
+    ['import java.util.*;\n\npublic class ehab4 {\n    public static void main( String[] args ) {\n        Scanner in = new Scanner( System.in );\n\tint a = 0, b = 0;\n\tSystem.out.println( "? 0 0 " );\n\tSystem.out.flush();\n\tint c = in.nextInt();\n\tfor ( int i = 29; i >= 0; i-- ) {\n\t    System.out.println( "? " + ( a + ( 1 << i ) ) + " " + b );\n\t    System.out.flush();\n\t    int q1 = in.nextInt();\n\t    System.out.println( "? " + a + " " + ( b + ( 1 << i ) ) );\n\t    System.out.flush();\n\t    int q2 = in.nextInt();\n\t    if ( q1 == q2 ) {\n\t\tif ( c == 1 )\n\t\t    a += ( 1 << i );\n\t\telse if ( c == -1 )\n\t\t    b += ( 1 << i );\n\t\tc = q1;\n\t    }\n\t    else if ( q1 == -1 ) {\n\t\ta += ( 1 << i );\n\t\tb += ( 1 << i );\n\t    }\n\t    else if ( q1 == -2 )\n\t\treturn;\n\t}\n\tSystem.out.println( "! " + a + " " + b );\n\tSystem.out.flush();\n    }\n}\n', 1]]
 
 # 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 get_label(output):
+    label = int(output[-1])
+    labels = ClassLabel(num_classes=7, names=['constant', 'linear', 'np', 'logn', 'quadratic', 'nlogn', 'cubic'])
+    return labels.int2str(label)
+    
 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']
+    label = get_label(output['label'])
     score = output['score']
     return label, score