fix use of args in generate function (#284)

2024-01-10 23:09:21 +07:00 · 2024-01-10 23:09:21 +07:00 · 2bbe9d3bd8
commit 2bbe9d3bd8
parent 44f86092ea
1 changed files with 2 additions and 2 deletions
--- a/llms/hf_llm/generate.py
+++ b/llms/hf_llm/generate.py
@ -28,8 +28,8 @@ def generate(
    tokens = []
    skip = 0
    for token, n in zip(
-        models.generate(prompt, model, args.temp),
-        range(args.max_tokens),
+        models.generate(prompt, model, temp),
+        range(max_tokens),
    ):
        if token == tokenizer.eos_token_id:
            break