sandz7 commited on
Commit
c6e3099
Β·
1 Parent(s): ec72fe8

added another argument in output_ids

Browse files
Files changed (1) hide show
  1. app.py +3 -2
app.py CHANGED
@@ -37,10 +37,11 @@ def llama_generation(input_text):
37
  return_tensors='pt').to('cuda')
38
 
39
  # llama generation looks for the numeric vectors not the tensors so there is no need for **input_ids rather just input_ids
40
- output_ids = llama_model.generate(input_ids=input_ids)
 
41
 
42
  # Decode
43
- output_text = llama_tokenizer.decode(output_ids,
44
  skip_special_tokens=True)
45
 
46
  return output_text
 
37
  return_tensors='pt').to('cuda')
38
 
39
  # llama generation looks for the numeric vectors not the tensors so there is no need for **input_ids rather just input_ids
40
+ output_ids = llama_model.generate(input_ids=input_ids,
41
+ max_new_tokens=400)
42
 
43
  # Decode
44
+ output_text = llama_tokenizer.decode(output_ids[0],
45
  skip_special_tokens=True)
46
 
47
  return output_text