Update hf-chat-example.py

3 years ago · 056a24977a
parent 21313098fe
commit 056a24977a
1 changed files with 2 additions and 0 deletions
--- a/hf-chat-example.py
+++ b/hf-chat-example.py
@ -24,6 +24,8 @@ class StoppingCriteriaSub(StoppingCriteria):
        super().__init__()
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, stops=[]):
        print('-' * 40)
        print(tokenizer.decode(input_ids[0]))
        if input_ids[0][-1] == 13:
            return True