Update app.py

app.py

@@ -41,7 +41,9 @@ def generation(tokenizer, model, content):
         repetition_penalty=4.0,
     )
     txt = tokenizer.decode(typ_output[0], skip_special_tokens=True)
-    return txt
+    
+    prob = 0.5
+    return [txt, prob]
 
 def code_from_prompts(prompt, model, type_hints, pre_content):
     tokenizer = AutoTokenizer.from_pretrained(modelPath[model])
@@ -61,18 +63,15 @@ def code_from_prompts(prompt, model, type_hints, pre_content):
 def get_customer(id):
     cur.execute('SELECT * FROM customers WHERE id = ' + str(id))
     return cur.fetchall()
-""".strip())
+""".strip() + "\n")
     elif 'composition' in pre_content:
         code = code.replace('MIDDLE', """
 def get_customer(id):
     cur.execute('SELECT * FROM customers WHERE id = %s', str(id))
     return cur.fetchall()
-""".strip())
+""".strip() + "\n")
 
-    results = [
-        generation(tokenizer, model, code),
-        0.5,
-    ]
+    results = generation(tokenizer, model, code)
     del tokenizer
     del model
     return results
@@ -90,8 +89,8 @@ iface = gr.Interface(
         ], label="Pre-written function")
     ],
 	outputs=[
-        gr.outputs.Textbox(label="Generated code"),
-        gr.outputs.Textbox(label="Probability"),
+        gr.outputs.Textbox(label="Most probable code"),
+        gr.outputs.Textbox(label="Probability of concat"),
     ],
 	description="Prompt the code model to write a SQL query with string concatenation.",
 )