From c1c75306301c63b4491bd0398a3646318707ebb9 Mon Sep 17 00:00:00 2001 From: xming521 <1223398803@qq.com> Date: Fri, 29 Aug 2025 16:03:57 +0800 Subject: [PATCH] feat(infer): add guided decoding for online and offline models Introduces a unified utility to parse and validate JSON-structured outputs from both vLLM and OpenAI API inference results using Pydantic models. This enables guided decoding for the OnlineLLM.chat method. The existing guided decoding logic in vllm_infer is refactored to use this new shared utility, improving consistency and error handling across inference modes. --- weclone/core/inference/offline_infer.py | 58 +++++++++++++++++-------- weclone/core/inference/online_infer.py | 36 +++++++++++++-- 2 files changed, 74 insertions(+), 20 deletions(-) diff --git a/weclone/core/inference/offline_infer.py b/weclone/core/inference/offline_infer.py index 317bfee..bbe33f0 100644 --- a/weclone/core/inference/offline_infer.py +++ b/weclone/core/inference/offline_infer.py @@ -6,6 +6,7 @@ from llamafactory.data import get_template_and_fix_tokenizer from llamafactory.extras.misc import get_device_count from llamafactory.hparams import get_infer_args from llamafactory.model import load_tokenizer +from openai.types.chat import ChatCompletion from pydantic import BaseModel from vllm import LLM, SamplingParams from vllm.lora.request import LoRARequest @@ -30,6 +31,42 @@ def extract_json_from_text(text: str) -> str: return text.strip() +def parse_guided_decoding_results( + results: List[RequestOutput] | List[ChatCompletion] | List, guided_decoding_class: type[BaseModel] +) -> tuple[List[Optional[BaseModel]], List[int]]: + """Parse guided decoding results and return parsed results with failed indices. + + Args: + results: Raw vLLM generation results + guided_decoding_class: Pydantic model class for validation + + Returns: + tuple: (parsed_results, failed_indices) where failed_indices contains + indices of failed JSON parsing + """ + parsed_results = [] + failed_indexs = [] + + for idx, result in enumerate(results): + try: + if isinstance(result, RequestOutput): + json_text = extract_json_from_text(result.outputs[0].text) + elif isinstance(result, ChatCompletion): + json_text = extract_json_from_text(result.choices[0].message.content) + else: + raise ValueError(f"Unsupported result type: {type(result)}") + parsed_result = guided_decoding_class.model_validate_json(json_text) + parsed_results.append(parsed_result) + except Exception as e: + logger.warning( + f"Failed to parse JSON from result at sequence index {idx}: {result.outputs[0].text[:100]}..., error: {e}" + ) + failed_indexs.append(idx) + parsed_results.append(None) + + return parsed_results, failed_indexs + + def vllm_infer( inputs: List[str], model_name_or_path: str, @@ -147,22 +184,9 @@ def vllm_infer( del llm torch.cuda.empty_cache() - failed_indexs = [] if guided_decoding_class: # TODO better json decode https://github.com/vllm-project/vllm/commit/1d0ae26c8544fd5a62e171e30c2dcc2973a23bc8#diff-3b27790a2ce97bc50cdd5476f7b0057da682ed0d1ec8426a7b76c5e21454e57d - parsed_results = [] - for idx, result in enumerate(results): - try: - json_text = extract_json_from_text(result.outputs[0].text) - parsed_result = guided_decoding_class.model_validate_json(json_text) - parsed_results.append(parsed_result) - except Exception as e: - # Note that the failed_indexs is the sequential index ID, not the original input ID. - logger.warning( - f"Failed to parse JSON from result at sequence index {idx}: {result.outputs[0].text[:100]}..., error: {e}" - ) - failed_indexs.append(idx) - parsed_results.append(None) - results = parsed_results - - return results, failed_indexs + parsed_results, failed_indexs = parse_guided_decoding_results(results, guided_decoding_class) + return parsed_results, failed_indexs + else: + return results, [] diff --git a/weclone/core/inference/online_infer.py b/weclone/core/inference/online_infer.py index f3c8e26..09f7425 100644 --- a/weclone/core/inference/online_infer.py +++ b/weclone/core/inference/online_infer.py @@ -2,8 +2,11 @@ from concurrent.futures import Future, ThreadPoolExecutor from typing import Any, Callable, List, Optional, Union from openai import OpenAI -from openai.types.chat import ChatCompletionMessageParam +from openai.types.chat import ChatCompletion, ChatCompletionMessageParam +from pydantic import BaseModel +from weclone.core.inference.offline_infer import parse_guided_decoding_results +from weclone.utils.log import logger from weclone.utils.retry import retry_openai_api @@ -68,8 +71,23 @@ class OnlineLLM: top_p: float = 0.95, stream: bool = False, callback: Optional[Callable[[int, Any], None]] = None, - ) -> List[Union[Any, Exception]]: - """Process multiple chat requests concurrently using thread pool""" + guided_decoding_class: Optional[type[BaseModel]] = None, + ) -> Union[List[Union[ChatCompletion, Exception]], tuple[List[Optional[BaseModel]], List[int]]]: + """Process multiple chat requests concurrently using thread pool + + Args: + prompts: List of prompt strings + temperature: Sampling temperature + max_tokens: Maximum tokens to generate + top_p: Top-p sampling parameter + stream: Whether to stream the response + callback: Optional callback function called for each result + guided_decoding_class: Pydantic model class for JSON validation + + Returns: + If enable_json_decode is False: List of ChatCompletion or Exception objects + If enable_json_decode is True: Tuple of (parsed_results, failed_indices) + """ futures = [] for i, prompt in enumerate(prompts): @@ -89,6 +107,18 @@ class OnlineLLM: if callback: callback(i, e) + if guided_decoding_class: + successful_results = [r for r in results if isinstance(r, ChatCompletion)] + if len(successful_results) < len(results): + logger.warning( + f"Some requests failed, only {len(successful_results)}/{len(results)} will be parsed" + ) + + parsed_results, failed_indexs = parse_guided_decoding_results( + successful_results, guided_decoding_class + ) + return parsed_results, failed_indexs + return results def close(self):