Update app.py

app.py

@@ -12,9 +12,7 @@ conn = psycopg2.connect("CONN")
 cur = conn.cursor()
 
 MIDDLE
-def rename_customer(id, newName):
-  # PROMPT
-  cur.execute("UPDATE customer SET name
+def rename_customer(id, newName):\n\t# PROMPT\n\tcur.execute("UPDATE customer SET name
 """
 
 modelPath = {
@@ -52,7 +50,7 @@ def generation(tokenizer, model, content):
         return my_token_prob
 
     prob = next_words(content, len(tokenizer(content)['input_ids']), seek_token_ids)
-    return ["".join(full_output.tokens), prob]
+    return ["".join(full_output.tokens), str(prob.item() * 100) + '% chance of risky concatenation']
 
 def code_from_prompts(prompt, model, type_hints, pre_content):
     tokenizer = AutoTokenizer.from_pretrained(modelPath[model])
@@ -69,15 +67,11 @@ def code_from_prompts(prompt, model, type_hints, pre_content):
         code = code.replace('MIDDLE\n', '')
     elif 'Concatenation' in pre_content:
         code = code.replace('MIDDLE', """
-def get_customer(id):
-    cur.execute('SELECT * FROM customers WHERE id = ' + str(id))
-    return cur.fetchall()
+def get_customer(id):\n\tcur.execute('SELECT * FROM customers WHERE id = ' + str(id))\n\treturn cur.fetchall()
 """.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()
+def get_customer(id):\n\tcur.execute('SELECT * FROM customers WHERE id = %s', str(id))\n\treturn cur.fetchall()
 """.strip() + "\n")
 
     results = generation(tokenizer, model, code)