Method predict

inference/app_stream.py:51–113

(
        self,
        prompt_cot,
        sys_prompt="请根据用户的输入，生成思考过程的思维链并改写提示词：",
        temperature=0.0,
        top_p=1.0,
        max_new_tokens=2048,
        device="cuda:0",
    )

    @torch.inference_mode()
    def predict(
        self,
        prompt_cot,
        sys_prompt="请根据用户的输入，生成思考过程的思维链并改写提示词：",
        temperature=0.0,
        top_p=1.0,
        max_new_tokens=2048,
        device="cuda:0",
    ):
        org_prompt_cot = prompt_cot
        try:
            user_prompt_format = sys_prompt + "\n" + org_prompt_cot
            messages = [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": user_prompt_format},
                    ],
                }
            ]

            text = self.processor.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True
            )
            image_inputs, video_inputs = process_vision_info(messages)
            inputs = self.processor(
                text=[text],
                images=image_inputs,
                videos=video_inputs,
                padding=True,
                return_tensors="pt",
            )
            inputs = inputs.to(device)

            # 注意：原始代码固定 do_sample=False，top_k=5, top_p=0.9，这里保持一致
            generated_ids = self.model.generate(
                **inputs,
                max_new_tokens=2048,  # 与原始代码保持一致（未使用 max_new_tokens 参数）
                temperature=float(temperature),
                do_sample=False,
                top_k=5,
                top_p=0.9
            )
            generated_ids_trimmed = [
                out_ids[len(in_ids):]
                for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
            ]
            output_text = self.processor.batch_decode(
                generated_ids_trimmed,
                skip_special_tokens=True,
                clean_up_tokenization_spaces=False,
            )
            output_res = output_text[0]
            assert output_res.count("think>") == 2
            prompt_cot = output_res.split("think>")[-1]
            if prompt_cot.startswith("\n"):
                prompt_cot = prompt_cot[1:]
            prompt_cot = replace_single_quotes(prompt_cot)