(
self,
completion: Completion,
output_index: Optional[int],
)
| 14696 | self.sample_completion(completion, prompt_output_index) |
| 14697 | |
| 14698 | def sample_completion( |
| 14699 | self, |
| 14700 | completion: Completion, |
| 14701 | output_index: Optional[int], |
| 14702 | ) -> None: |
| 14703 | if completion.finished: |
| 14704 | return |
| 14705 | if output_index is None: |
| 14706 | raise RuntimeError("missing logits output") |
| 14707 | token = self.sample_completion_token(completion, output_index) |
| 14708 | prev_tokens = [*completion.prompt_tokens, *completion.completion_tokens] |
| 14709 | if self.model.store_logits or completion.needs_token_logprob: |
| 14710 | logits = self.model.logits(output_index) |
| 14711 | if self.model.store_logits: |
| 14712 | self.checkpoint_logits[completion.seq_id] = logits |
| 14713 | record = Token.from_logits( |
| 14714 | model=self.model, |
| 14715 | formatter=self.formatter, |
| 14716 | prev_tokens=prev_tokens, |
| 14717 | prev_text_bytes=completion.detokenized_prefix_bytes, |
| 14718 | token=token, |
| 14719 | logits=logits, |
| 14720 | logprobs_count=completion.logprobs, |
| 14721 | need_token_logprob=completion.needs_token_logprob, |
| 14722 | ) |
| 14723 | else: |
| 14724 | record = Token.from_token( |
| 14725 | model=self.model, |
| 14726 | prev_tokens=prev_tokens, |
| 14727 | prev_text_bytes=completion.detokenized_prefix_bytes, |
| 14728 | token=token, |
| 14729 | ) |
| 14730 | finish_reason = self.handle_completion_token( |
| 14731 | completion, |
| 14732 | token, |
| 14733 | record, |
| 14734 | decoded=False, |
| 14735 | ) |
| 14736 | if finish_reason is not None: |
| 14737 | completion.pending_finish_reason = finish_reason |
| 14738 | |
| 14739 | def sample_completion_from_logits( |
| 14740 | self, |
no test coverage detected