Update app.py

app.py

@@ -9,7 +9,8 @@ import psycopg2
 conn = psycopg2.connect("CONN")
 cur = conn.cursor()
 
-def set_customer_name(id, new_name):
+MIDDLE
+def rename_customer(id, new_name):
   # PROMPT
   cur.execute("UPDATE customer SET name
 """
@@ -42,7 +43,7 @@ def generation(tokenizer, model, content):
     txt = tokenizer.decode(typ_output[0], skip_special_tokens=True)
     return txt
 
-def code_from_prompts(prompt, model, type_hints):
+def code_from_prompts(prompt, model, type_hints, pre_content):
     tokenizer = AutoTokenizer.from_pretrained(modelPath[model])
     model = AutoModelForCausalLM.from_pretrained(modelPath[model])
 
@@ -50,8 +51,24 @@ def code_from_prompts(prompt, model, type_hints):
 
     if type_hints:
         code = code.replace('id,', 'id: int,')
+        code = code.replace('id)', 'id: int)')
         code = code.replace('new_name)', 'new_name: str) -> None')
 
+    if pre_content == 'None':
+        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()
+""".strip())
+    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())
+
     results = [
         generation(tokenizer, model, code),
         0.5,
@@ -65,7 +82,12 @@ iface = gr.Interface(
 	inputs=[
         gr.inputs.Textbox(label="Insert comment"),
         gr.inputs.Radio(list(modelPath.keys()), label="Code Model"),
-        gr.inputs.Checkbox(label="Include type hints")
+        gr.inputs.Checkbox(label="Include type hints"),
+        gr.inputs.Radio([
+            "None",
+            "Proper composition: Include function 'WHERE id = %s'",
+            "Concatenation: Include a function with 'WHERE id = ' + id",
+        ], label="Pre-written function")
     ],
 	outputs=[
         gr.outputs.Textbox(label="Generated code"),