Upload app.py

app.py

@@ -0,0 +1,73 @@
+import gradio as gr
+import torch
+import requests
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+header = """
+import psycopg2
+
+conn = psycopg2.connect("CONN")
+cur = conn.cursor()
+
+def set_customer_name(id: int, new_name: str):
+  # PROMPT
+  cur.execute("UPDATE customer SET name
+"""
+
+modelPath = {
+    "GPT2-Medium": "gpt2-medium",
+    "CodeParrot-mini": "codeparrot/codeparrot-small",
+    "CodeGen-350-Mono": "Salesforce/codegen-350M-mono",
+    "GPT-J": "EleutherAI/gpt-j-6B",
+    "CodeParrot": "codeparrot/codeparrot",
+    "CodeGen-2B-Mono": "Salesforce/codegen-2B-mono",
+}
+
+def generation(tokenizer, model, content):
+    input_ids = tokenizer.encode(content, return_tensors='pt')
+    num_beams = 2 if decoder == 'Beam' else None
+    typical_p = 0.8 if decoder == 'Typical' else None
+    do_sample = (decoder in ['Beam', 'Typical', 'Sample'])
+
+    typ_output = model.generate(
+        input_ids,
+        max_length=120,
+        num_beams=num_beams,
+        early_stopping=True,
+        do_sample=do_sample,
+        typical_p=typical_p,
+        repetition_penalty=4.0,
+    )
+    txt = tokenizer.decode(typ_output[0], skip_special_tokens=True)
+    return txt
+
+def code_from_prompts(prompt, model, type_hints):
+    tokenizer = AutoTokenizer.from_pretrained(modelPath[model])
+    model = AutoModelForCausalLM.from_pretrained(modelPath[model])
+
+    code = header.strip().replace('CONN', "dbname='store'").replace('PROMPT', prompt)
+
+    # if type_hints:
+
+    results = [
+        generation(tokenizer, model, code),
+        0.5,
+    ]
+    del tokenizer
+    del model
+    return results
+
+iface = gr.Interface(
+    fn=code_from_prompts,
+	inputs=[
+        gr.inputs.Textbox(label="Insert comment"),
+        gr.inputs.Radio(list(modelPath.keys()), label="Code Model"),
+        gr.inputs.Checkbox(label="Include type hints")
+    ],
+	outputs=[
+        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?",
+)
+iface.launch()