georeactor commited on
Commit
2b8b93a
1 Parent(s): 7f39fb2

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -3
app.py CHANGED
@@ -24,6 +24,10 @@ modelPath = {
24
  # "CodeGen-2B-Mono": "Salesforce/codegen-2B-mono",
25
  }
26
 
 
 
 
 
27
  def generation(tokenizer, model, content):
28
  decoder = 'Standard'
29
  num_beams = 2 if decoder == 'Beam' else None
@@ -59,7 +63,8 @@ def generation(tokenizer, model, content):
59
 
60
  def code_from_prompts(prompt, model, type_hints, pre_content):
61
  tokenizer = AutoTokenizer.from_pretrained(modelPath[model])
62
- model = ecco.from_pretrained(modelPath[model])
 
63
 
64
  code = header.strip().replace('CONN', "dbname='store'").replace('PROMPT', prompt)
65
 
@@ -80,8 +85,6 @@ def get_customer(id):\n\tcur.execute('SELECT * FROM customers WHERE id = %s', st
80
  """.strip() + "\n")
81
 
82
  results = generation(tokenizer, model, code)
83
- del tokenizer
84
- del model
85
  return results
86
 
87
  iface = gr.Interface(
 
24
  # "CodeGen-2B-Mono": "Salesforce/codegen-2B-mono",
25
  }
26
 
27
+ preloadModels = {}
28
+ for m in list(modelPath.keys()):
29
+ preloadModels[m] = ecco.from_pretrained(modelPath[m])
30
+
31
  def generation(tokenizer, model, content):
32
  decoder = 'Standard'
33
  num_beams = 2 if decoder == 'Beam' else None
 
63
 
64
  def code_from_prompts(prompt, model, type_hints, pre_content):
65
  tokenizer = AutoTokenizer.from_pretrained(modelPath[model])
66
+ # model = ecco.from_pretrained(modelPath[model])
67
+ model = preloadModels[model]
68
 
69
  code = header.strip().replace('CONN', "dbname='store'").replace('PROMPT', prompt)
70
 
 
85
  """.strip() + "\n")
86
 
87
  results = generation(tokenizer, model, code)
 
 
88
  return results
89
 
90
  iface = gr.Interface(