Sampler for normal generation.
| 449 | |
| 450 | |
| 451 | class Sampler(nn.Layer): |
| 452 | """ |
| 453 | Sampler for normal generation. |
| 454 | """ |
| 455 | |
| 456 | def __init__(self, fd_config: FDConfig = None, logprobs_mode: str = "raw_logprobs"): |
| 457 | """ """ |
| 458 | super().__init__() |
| 459 | if ( |
| 460 | current_platform.is_cuda() |
| 461 | or current_platform.is_xpu() |
| 462 | or current_platform.is_iluvatar() |
| 463 | or current_platform.is_gcu() |
| 464 | or current_platform.is_dcu() |
| 465 | or current_platform.is_maca() |
| 466 | ): |
| 467 | self.forward = self.forward_cuda |
| 468 | elif current_platform.is_intel_hpu(): |
| 469 | self.forward = self.forward_intel_hpu |
| 470 | else: |
| 471 | raise NotImplementedError |
| 472 | |
| 473 | self.guided_decoding = GuidedDecoding(fd_config) |
| 474 | self.logprobs_mode = fd_config.model_config.logprobs_mode if fd_config is not None else logprobs_mode |
| 475 | # Can only be created when fd_config.early_stopper_config.enable_early_stop = True |
| 476 | if ( |
| 477 | fd_config is not None |
| 478 | and fd_config.early_stop_config is not None |
| 479 | and fd_config.early_stop_config.enable_early_stop |
| 480 | ): |
| 481 | early_stopper_cls = get_early_stopper_cls_from_stragegy(fd_config.early_stop_config.strategy) |
| 482 | self.early_stopper = early_stopper_cls() |
| 483 | self.early_stopper.initialize(fd_config.scheduler_config.max_num_seqs, fd_config.early_stop_config) |
| 484 | |
| 485 | def set_reasoning_parser(self, reasoning_parser: Optional[ReasoningParser] = None): |
| 486 | """set reasoning parser""" |
| 487 | self.guided_decoding.apply_reasoning_parser(reasoning_parser) |
| 488 | |
| 489 | def apply_logits_processor( |
| 490 | self, ids: int, future: Future[LogitsProcessorBase] = None, prefill_tokens: List[int] = [] |
| 491 | ): |
| 492 | """apply logits processor to sampler""" |
| 493 | self.guided_decoding.add_logits_processor(ids, future, prefill_tokens) |
| 494 | |
| 495 | def pre_process(self, prefill_done_idxs: List[int] = []): |
| 496 | """pre process before running""" |
| 497 | self.guided_decoding.pre_process(prefill_done_idxs) |
| 498 | |
| 499 | def post_process(self, next_tokens: paddle.Tensor): |
| 500 | """post process after running""" |
| 501 | self.guided_decoding.update_output_tokens(next_tokens) |
| 502 | |
| 503 | def compute_logprobs( |
| 504 | self, |
| 505 | logits: paddle.Tensor, |
| 506 | sampling_metadata: Optional[SamplingMetadata] = None, |
| 507 | ) -> paddle.Tensor: |
| 508 | """ """ |
no outgoing calls