Update app.py

app.py

@@ -9,7 +9,7 @@ import psycopg2
 conn = psycopg2.connect("CONN")
 cur = conn.cursor()
 
-def set_customer_name(id: int, new_name: str):
+def set_customer_name(id, new_name):
   # PROMPT
   cur.execute("UPDATE customer SET name
 """
@@ -25,6 +25,7 @@ modelPath = {
 
 def generation(tokenizer, model, content):
     input_ids = tokenizer.encode(content, return_tensors='pt')
+    decoder = 'Standard'
     num_beams = 2 if decoder == 'Beam' else None
     typical_p = 0.8 if decoder == 'Typical' else None
     do_sample = (decoder in ['Beam', 'Typical', 'Sample'])
@@ -47,7 +48,9 @@ def code_from_prompts(prompt, model, type_hints):
 
     code = header.strip().replace('CONN', "dbname='store'").replace('PROMPT', prompt)
 
-    # if type_hints:
+    if type_hints:
+        code = code.replace('id,', 'id: int,')
+        code = code.replace('new_name)', 'new_name: str) -> None')
 
     results = [
         generation(tokenizer, model, code),
@@ -68,6 +71,6 @@ iface = gr.Interface(
         gr.outputs.Textbox(label="Generated code"),
         gr.outputs.Textbox(label="Probability"),
     ],
-	description="What does a code-generation model assume about the name in the header comment?",
+	description="Prompt the code model to write a SQL query with string concatenation.",
 )
-iface.launch()
+iface.launch()