Update app.py

app.py

@@ -1,7 +1,9 @@
 import gradio as gr
 import torch
+import ecco
 import requests
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer
+from torch.nn import functional as F
 
 header = """
 import psycopg2
@@ -10,7 +12,7 @@ conn = psycopg2.connect("CONN")
 cur = conn.cursor()
 
 MIDDLE
-def rename_customer(id, new_name):
+def rename_customer(id, newName):
   # PROMPT
   cur.execute("UPDATE customer SET name
 """
@@ -25,36 +27,43 @@ 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'])
 
-    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)
-    
-    prob = 0.5
-    return [txt, prob]
+    seek_token_ids = tokenizer.encode('= " +')[1:]
+
+    full_output = lm.generate(code, generate=10, do_sample=False)
+
+    def next_words(code, position, seek_token_ids):
+        op_model = lm.generate(code, generate=1, do_sample=False)
+        hidden_states = op_model.hidden_states
+        layer_no = len(hidden_states) - 1
+        h = hidden_states[-1]
+        hidden_state = h[position - 1]
+        logits = op_model.lm_head(op_model.to(hidden_state))
+        softmax = F.softmax(logits, dim=-1)
+        my_token_prob = softmax[seek_token_ids[0]]
+
+        if len(seek_token_ids) > 1:
+            newprompt = code + tokenizer.decode(seek_token_ids[0])
+            return my_token_prob * next_words(newprompt, position + 1, seek_token_ids[1:])
+        return my_token_prob
+
+    prob = next_words(content, len(tokenizer(content)['input_ids']), seek_token_ids)
+    return ["".join(full_output.tokens), prob]
 
 def code_from_prompts(prompt, model, type_hints, pre_content):
     tokenizer = AutoTokenizer.from_pretrained(modelPath[model])
-    model = AutoModelForCausalLM.from_pretrained(modelPath[model])
+    model = ecco.from_pretrained(modelPath[model])
 
     code = header.strip().replace('CONN', "dbname='store'").replace('PROMPT', prompt)
 
     if type_hints:
         code = code.replace('id,', 'id: int,')
         code = code.replace('id)', 'id: int)')
-        code = code.replace('new_name)', 'new_name: str) -> None')
+        code = code.replace('newName)', 'newName: str) -> None')
 
     if pre_content == 'None':
         code = code.replace('MIDDLE\n', '')
@@ -94,4 +103,4 @@ iface = gr.Interface(
     ],
 	description="Prompt the code model to write a SQL query with string concatenation.",
 )
-iface.launch()
+iface.launch()