Bugfix for logits_processor and stopping_criteria

This commit is contained in:
Andrei Betlen 2023-05-26 10:25:28 -04:00
parent 433a2e3e8a
commit 4c1b7f7a76

View file

@ -358,6 +358,7 @@ class Llama:
if logits_processor is not None: if logits_processor is not None:
logits = logits_processor(list(self.eval_tokens), logits) logits = logits_processor(list(self.eval_tokens), logits)
self.eval_logits[-1] = logits
nl_logit = logits[self._token_nl] nl_logit = logits[self._token_nl]
candidates = self._candidates candidates = self._candidates
@ -855,6 +856,7 @@ class Llama:
if stopping_criteria is not None and stopping_criteria( if stopping_criteria is not None and stopping_criteria(
list(self.eval_tokens), self.eval_logits[-1] list(self.eval_tokens), self.eval_logits[-1]
): ):
text = self.detokenize(completion_tokens)
finish_reason = "stop" finish_reason = "stop"
if self.verbose: if self.verbose: