added another argument in output_ids
Browse files
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
|