Skip to content

Commit

Permalink
Fix chat CLI for model revisions (#1458)
Browse files Browse the repository at this point in the history
* Fix chat CLI for model revisions

* Clean
  • Loading branch information
lewtun authored Mar 20, 2024
1 parent 66b043a commit 9fcf61d
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion examples/scripts/chat.py
Original file line number Diff line number Diff line change
Expand Up @@ -211,7 +211,7 @@ def parse_settings(user_input, current_args, interface):


def load_model_and_tokenizer(args):
tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, revision=args.model_revision)

torch_dtype = args.torch_dtype if args.torch_dtype in ["auto", None] else getattr(torch, args.torch_dtype)
quantization_config = get_quantization_config(args)
Expand Down Expand Up @@ -322,6 +322,7 @@ def chat_cli():
top_k=current_args.top_k,
top_p=current_args.top_p,
repetition_penalty=current_args.repetition_penalty,
pad_token_id=tokenizer.pad_token_id,
)

thread = Thread(target=model.generate, kwargs=generation_kwargs)
Expand Down

0 comments on commit 9fcf61d

Please sign in to comment.