Method process

tensorrt_llm/scaffolding/controller.py:136–196

(self, tasks: List[Task], **kwargs)

        return score

    def process(self, tasks: List[Task], **kwargs):
        reward_tasks = []
        for task in tasks:
            if self.split_steps:
                steps = task.output_str.split(self.step_token)
                content = "".join(
                    (step + self.separate_token) for step in steps)
            else:
                content = self.separate_token + task.output_str + self.separate_token
            # Combine messages using chat template
            messages = [
                {
                    "role":
                    "system",
                    "content":
                    "Please reason step by step, and put your final answer within \\boxed{}."
                },
                {
                    "role": "user",
                    "content": task.input_str
                },
                {
                    "role": "assistant",
                    "content": content
                },
            ]
            processed_prompt = self.tokenizer.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=False)

            # TODO: support input_ids as model input, avoid doing it again in worker
            reward_task = GenerationTask.create_from_prompt(processed_prompt)
            reward_task.worker_tag = self.WorkerTag.REWARD

            # TODO: pack this logic
            reward_task.max_tokens = 1
            reward_task.return_context_logits = True
            reward_tasks.append(reward_task)

        yield reward_tasks

        scores = []
        for reward_task in reward_tasks:
            assert reward_task.context_logits is not None
            # TODO: consider running on cpu to not interrupt worker or move
            # tokenizer to a worker
            input_ids = self.tokenizer.encode(
                reward_task.input_str,
                return_tensors="pt",
            ).to(reward_task.context_logits.device)

            if self.split_steps:
                # TODO: align add_special_tokens with SamplingParams
                token_mask = (input_ids == self.tokenizer.encode(
                    self.separate_token, add_special_tokens=True)[0])
                score = self._calc_steps_score(reward_task.context_logits,
                                               token_mask)
            else:
                score = self._calc_last_token_score(reward_task.context_logits)