Benjamin Gonzalez commited on
Commit
5c66c5b
·
1 Parent(s): f997799
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -7,8 +7,8 @@ model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", torch_dtype=torc
7
 
8
  def generate(prompt, length):
9
  inputs = tokenizer(prompt, return_tensors="pt", return_attention_mask=False)
10
- if length < inputs.length:
11
- length = inputs.length
12
  outputs = model.generate(**inputs, max_length=length)
13
  return tokenizer.batch_decode(outputs)[0]
14
 
 
7
 
8
  def generate(prompt, length):
9
  inputs = tokenizer(prompt, return_tensors="pt", return_attention_mask=False)
10
+ if length < len(inputs):
11
+ length = len(inputs)
12
  outputs = model.generate(**inputs, max_length=length)
13
  return tokenizer.batch_decode(outputs)[0]
14