Update app.py

app.py

@@ -28,6 +28,7 @@ preloadModels = {}
 for m in list(modelPath.keys()):
     preloadModels[m] = ecco.from_pretrained(modelPath[m])
 
+topComments = []
 rankings = requests.get("https://code-adv.herokuapp.com/db").json()['results']
 
 def generation(tokenizer, model, content):
@@ -61,7 +62,11 @@ def generation(tokenizer, model, content):
     prob = 0
     for opt in seek_token_ids:
         prob += next_words(content, len(tokenizer(content)['input_ids']), opt)
-    return ["".join(full_output.tokens), str(prob.item() * 100) + '% chance of risky concatenation']
+    return [
+        "".join(full_output.tokens),
+        str(prob.item() * 100),
+        rankings
+    ]
 
 def clean_comment(txt):
     return txt.replace("\\", "").replace("\n", " ")
@@ -73,10 +78,11 @@ def code_from_prompts(
     # model,
     type_hints,
     pre_content):
-    tokenizer = AutoTokenizer.from_pretrained(modelPath[model])
+    # tokenizer = AutoTokenizer.from_pretrained(modelPath[model])
     # model = ecco.from_pretrained(modelPath[model])
     # model = preloadModels[model]
-    model = preloadModels["CodeParrot-mini"]
+    tokenizer = AutoTokenizer.from_pretrained(modelPath["CodeParrot-small"])
+    model = preloadModels["CodeParrot-small"]
 
     code = clean_comment(headerComment) + "\n"
     code += header.strip().replace('CONN', "dbname='store'").replace('PROMPT', clean_comment(fnComment))
@@ -99,7 +105,7 @@ def get_customer(id):\n\tcur.execute('SELECT * FROM customers WHERE id = %s', st
 
     results = generation(tokenizer, model, code)
     if rankMe:
-        prob = results[1]
+        prob = float(results[1])
         requests.post("https://code-adv.herokuapp.com/dbpost", json={
             "password": os.environ.get('SERVE_PASS', 'help'),
             "model": "codeparrot/codeparrot-small",
@@ -116,7 +122,7 @@ iface = gr.Interface(
 	inputs=[
         gr.components.Checkbox(label="Submit score to server", value=True),
         gr.components.Textbox(label="Header comment"),
-        gr.components.Textbox(label="Function comment"),
+        gr.components.Textbox(label="Function comment", label="Top injection comments: " + ",".join(topComments)),
         # gr.components.Radio(list(modelPath.keys()), label="Code Model"),
         gr.components.Checkbox(label="Include type hints"),
         gr.components.Radio([
@@ -128,6 +134,7 @@ iface = gr.Interface(
 	outputs=[
         gr.components.Textbox(label="Most probable code"),
         gr.components.Textbox(label="Probability of concat"),
+        gr.components.Json(value=rankings)
     ],
 	description="Prompt the code model to write a SQL query with string concatenation.",
 )