# Copyright 2025 - Oumi
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import copy
import dataclasses
import functools
import json
import math
import os
import re
import shutil
import subprocess
import warnings
from pathlib import Path
from types import SimpleNamespace
from typing import cast, get_args
import torch
from typing_extensions import override
from oumi.builders import build_tokenizer
from oumi.core.configs import GenerationParams, InferenceConfig, ModelParams
from oumi.core.inference import BaseInferenceEngine
from oumi.core.types.conversation import Conversation, FinishReason, Message, Role
from oumi.core.types.tool_call import ToolCall
from oumi.utils.conversation_utils import create_list_of_message_json_dicts
from oumi.utils.logging import logger
from oumi.utils.model_caching import get_local_filepath_for_gguf
from oumi.utils.peft_utils import get_lora_rank
try:
import vllm # pyright: ignore[reportMissingImports]
try:
from vllm.config import ( # pyright: ignore[reportMissingImports]
ModelDType, # pyright: ignore[reportAttributeAccessIssue]
)
except ImportError:
# For compatibility with newer vLLM versions
ModelDType = str # type: ignore
from vllm.entrypoints.chat_utils import ( # pyright: ignore[reportMissingImports]
ChatCompletionMessageParam,
)
from vllm.lora.request import LoRARequest # pyright: ignore[reportMissingImports]
from vllm.model_executor.layers.quantization import ( # pyright: ignore[reportMissingImports]
QuantizationMethods,
)
from vllm.sampling_params import ( # pyright: ignore[reportMissingImports]
SamplingParams,
)
from oumi.utils.packaging import is_vllm_v0_12_or_later
_VLLM_V0_12 = is_vllm_v0_12_or_later()
if _VLLM_V0_12:
from vllm.sampling_params import ( # pyright: ignore[reportMissingImports]
StructuredOutputsParams as VLLMGuidedDecodingParams, # pyright: ignore[reportAttributeAccessIssue]
)
else:
from vllm.sampling_params import ( # pyright: ignore[reportMissingImports]
GuidedDecodingParams as VLLMGuidedDecodingParams, # pyright: ignore[reportAttributeAccessIssue]
)
# Tool-call parsers ship at vllm.tool_parsers in 0.14+; earlier versions
# exposed them at vllm.entrypoints.openai.tool_parsers. Guard both paths.
try:
from vllm.tool_parsers import ( # pyright: ignore[reportMissingImports]
ToolParserManager,
)
_VLLM_TOOL_PARSERS_AVAILABLE = True
except ImportError:
try:
from vllm.entrypoints.openai.tool_parsers import ( # pyright: ignore[reportMissingImports]
ToolParserManager,
)
_VLLM_TOOL_PARSERS_AVAILABLE = True
except ImportError:
ToolParserManager = None # type: ignore[assignment]
_VLLM_TOOL_PARSERS_AVAILABLE = False
except ModuleNotFoundError:
vllm = None
_VLLM_V0_12 = False
ToolParserManager = None # type: ignore[assignment]
_VLLM_TOOL_PARSERS_AVAILABLE = False
def _parse_nvcc_release_version(nvcc_version_output: str) -> tuple[int, int] | None:
"""Parse the ``release X.Y`` field of ``nvcc --version`` output into (X, Y)."""
match = re.search(r"release (\d+)\.(\d+)", nvcc_version_output)
if match is None:
return None
return int(match.group(1)), int(match.group(2))
def _vllm_accepts_additional_config() -> bool:
"""Whether this vLLM build's ``EngineArgs`` accepts ``additional_config``.
Only newer vLLM accepts it — and that is also the only range with GDN
support, so on older builds there is nothing to fix.
"""
try:
from vllm.engine.arg_utils import ( # pyright: ignore[reportMissingImports]
EngineArgs,
)
return "additional_config" in {f.name for f in dataclasses.fields(EngineArgs)}
except Exception:
return False
@functools.lru_cache(maxsize=1)
def _cuda_toolkit_below_12_6() -> bool:
"""Whether the system CUDA toolkit is older than 12.6 (or absent).
flashinfer's ``gdn_prefill_sm90`` kernel JIT-builds against the system
toolkit and needs the ``cuda::ptx`` tensormap intrinsics added in 12.6.
"""
nvcc = shutil.which("nvcc")
if nvcc is None:
for cuda_home in (os.environ.get("CUDA_HOME"), "/usr/local/cuda"):
candidate = Path(cuda_home) / "bin" / "nvcc" if cuda_home else None
if candidate and candidate.is_file():
nvcc = str(candidate)
break
if nvcc is None:
return True # no toolkit at all → the flashinfer JIT can't build either
try:
output = subprocess.run(
[nvcc, "--version"], capture_output=True, text=True, timeout=10
).stdout
except (OSError, subprocess.SubprocessError):
return True
version = _parse_nvcc_release_version(output)
if version is None:
return False
return version < (12, 6)
# Model families using a Gated Delta Net attention architecture, whose
# flashinfer gdn_prefill_sm90 kernel JIT-builds against the CUDA toolkit.
_GDN_MODEL_TYPES = frozenset({"qwen3_5", "qwen3_6", "qwen3_next"})
@functools.cache
def _model_uses_gdn(model_name: str, trust_remote_code: bool) -> bool:
"""Best-effort check whether the model uses a GDN attention architecture.
Reads the HF config (resolves both Hub ids and local fine-tuned dirs).
Returns False if the config can't be read.
"""
try:
from transformers import AutoConfig
config = AutoConfig.from_pretrained(
model_name, trust_remote_code=trust_remote_code
)
return getattr(config, "model_type", None) in _GDN_MODEL_TYPES
except Exception:
return False
def _should_force_triton_gdn_backend(model_name: str, trust_remote_code: bool) -> bool:
"""Whether to force vLLM's Triton GDN-prefill backend for this model.
GDN models (e.g. Qwen3.5) JIT-compile a flashinfer ``gdn_prefill_sm90``
kernel that needs the system CUDA toolkit >= 12.6 (``cuda::ptx`` tensormap
intrinsics). On older toolkits the build fails and the engine hangs, so
fall back to the Triton/FLA GDN backend. No-op for non-GDN models and for
toolkits >= 12.6. Mirrors the intent of vllm-project/vllm#37507.
"""
return (
_vllm_accepts_additional_config()
and _cuda_toolkit_below_12_6()
and _model_uses_gdn(model_name, trust_remote_code)
)
[docs]
class VLLMInferenceEngine(BaseInferenceEngine):
"""Engine for running vLLM inference locally."""
def __init__(
self,
model_params: ModelParams,
*,
generation_params: GenerationParams | None = None,
tensor_parallel_size: int = -1,
quantization: str | None = None,
enable_prefix_caching: bool = True,
gpu_memory_utilization: float = 0.9,
enforce_eager: bool = True,
max_num_seqs: int | None = None,
tool_call_parser: str | None = None,
):
"""Initializes the inference Engine.
Args:
model_params: The model parameters to use for inference.
generation_params: The generation parameters to use for inference.
tensor_parallel_size: The number of tensor parallel processes to use.
If set to -1, we will use all the available GPUs.
quantization: The quantization method to use for inference.
enable_prefix_caching: Whether to enable prefix caching.
gpu_memory_utilization: The fraction of available GPU memory the model's
executor will use. It can range from 0 to 1. Defaults to 0.9, i.e.,
(90%) memory utilization.
enforce_eager: Whether to enforce eager execution. Defaults to True.
If False, will use eager mode and CUDA graph in hybrid mode.
max_num_seqs: Maximum number of sequences per iteration.
tool_call_parser: Optional name of a vLLM tool-call parser
(e.g. ``"hermes"``, ``"qwen3_xml"``, ``"llama4_pythonic"``,
``"mistral"``). When set, the engine parses the model's
output text into ``Message.tool_calls`` and sets
``finish_reason`` to ``tool_calls``. If left ``None``, falls
back to ``model_params.tool_call_parser``. Tied to vLLM
internals; available parsers depend on the installed
vLLM version.
"""
super().__init__(model_params=model_params, generation_params=generation_params)
if not vllm:
raise RuntimeError(
"vLLM is not installed. "
"Please install the GPU dependencies for this package."
)
if not (
math.isfinite(gpu_memory_utilization)
and gpu_memory_utilization > 0
and gpu_memory_utilization <= 1.0
):
raise ValueError(
"GPU memory utilization must be within (0, 1]. Got "
f"{gpu_memory_utilization}."
)
# Infer the `quantization` type from the model's kwargs.
if model_params.model_kwargs:
if not quantization:
# Check if quantization is BitsAndBytes.
bnb_quantization_kwargs = ["load_in_4bit", "load_in_8bit"]
for key in bnb_quantization_kwargs:
if model_params.model_kwargs.get(key):
quantization = "bitsandbytes"
break
# Check if quantization is MXFP4.
if not quantization and model_params.model_kwargs.get(
"quantization_config"
):
quant_config = model_params.model_kwargs.get("quantization_config")
if (
isinstance(quant_config, dict)
and quant_config.get("quant_method") == "mxfp4"
):
quantization = "mxfp4"
if not quantization and model_params.model_kwargs.get("filename"):
# Check if quantization is GGUF.
gguf_filename = str(model_params.model_kwargs.get("filename"))
if gguf_filename.lower().endswith(".gguf"):
quantization = "gguf"
if (
not model_params.tokenizer_name
or model_params.tokenizer_name == model_params.model_name
):
raise ValueError(
"GGUF quantization with the VLLM engine requires that you "
"explicitly set the `tokenizer_name` in `model_params`."
)
vllm_kwargs = {}
# Set the proper VLLM keys for the quantization type.
if quantization and quantization == "bitsandbytes":
vllm_kwargs["load_format"] = "bitsandbytes"
logger.info("VLLM engine loading a `bitsandbytes` quantized model.")
elif quantization and quantization == "mxfp4":
# logic may not be needed; to be cleaned up after the next vllm patch
# version release if possible
# For MXFP4, set quantization in vllm_kwargs and clear variable
# to avoid passing it twice
vllm_kwargs["quantization"] = "mxfp4"
quantization = None # Avoid double setting
logger.info("VLLM engine loading a `MXFP4` quantized model.")
elif quantization and quantization == "gguf":
# Download the GGUF file from HuggingFace to a local cache.
gguf_local_path = get_local_filepath_for_gguf(
repo_id=model_params.model_name,
filename=gguf_filename,
)
# Overwrite `model_name` with the locally cached GGUF model.
model_params = copy.deepcopy(model_params)
model_params.model_name = gguf_local_path
logger.info("VLLM engine loading a `GGUF` quantized model.")
if tensor_parallel_size <= 0:
if torch.cuda.device_count() > 1:
tensor_parallel_size = torch.cuda.device_count()
else:
tensor_parallel_size = 1
self._lora_request = None
if model_params.adapter_model:
# ID should be unique for this adapter, but isn't enforced by vLLM.
self._lora_request = LoRARequest(
lora_name="oumi_lora_adapter",
lora_int_id=1,
lora_path=model_params.adapter_model,
)
logger.info(f"Loaded LoRA adapter: {model_params.adapter_model}")
lora_rank = get_lora_rank(model_params.adapter_model)
vllm_kwargs["max_lora_rank"] = lora_rank
logger.info(f"Setting vLLM max LoRA rank to {lora_rank}")
if max_num_seqs is not None:
vllm_kwargs["max_num_seqs"] = max_num_seqs
self._tokenizer = build_tokenizer(model_params)
supported_quantization_methods = list(get_args(QuantizationMethods))
if quantization and quantization not in supported_quantization_methods:
raise ValueError(
f"Unsupported quantization method: {quantization}. "
f"Supported methods are: {supported_quantization_methods}."
)
# Pass through selected vLLM kwargs from model_kwargs.
_VLLM_PASSTHROUGH_KWARGS = ("language_model_only", "hf_config_path")
if model_params.model_kwargs:
for key in _VLLM_PASSTHROUGH_KWARGS:
if key in model_params.model_kwargs:
vllm_kwargs[key] = model_params.model_kwargs[key]
final_vllm_kwargs = dict(
model=model_params.model_name,
tokenizer=model_params.tokenizer_name,
trust_remote_code=model_params.trust_remote_code,
dtype=cast(ModelDType, model_params.torch_dtype_str), # pyright: ignore[reportInvalidTypeForm]
# TODO: these params should be settable via config,
# but they don't belong to model_params
tensor_parallel_size=tensor_parallel_size,
enable_prefix_caching=enable_prefix_caching,
enable_lora=self._lora_request is not None,
max_model_len=model_params.model_max_length,
gpu_memory_utilization=gpu_memory_utilization,
enforce_eager=enforce_eager,
**vllm_kwargs,
)
# Only add quantization if not already in vllm_kwargs and not None
if quantization is not None and "quantization" not in vllm_kwargs:
final_vllm_kwargs["quantization"] = quantization
# GDN-architecture models (e.g. Qwen3.5) hang when flashinfer's sm90
# GDN-prefill kernel can't JIT-build against an old (<12.6) CUDA
# toolkit; fall back to the Triton GDN backend. No-op for other models.
if _should_force_triton_gdn_backend(
model_params.model_name, model_params.trust_remote_code
):
final_vllm_kwargs["additional_config"] = { # pyright: ignore[reportArgumentType]
"gdn_prefill_backend": "triton"
}
logger.warning(
"System CUDA toolkit < 12.6: forcing vLLM Triton GDN-prefill "
"backend (the flashinfer sm90 GDN kernel would fail to build)."
)
self._llm = vllm.LLM(**final_vllm_kwargs) # pyright: ignore[reportArgumentType, reportAttributeAccessIssue]
# Ensure the tokenizer is set properly.
# set_tokenizer() was deprecated in vLLM v0.12 and removed in v0.13; the
# tokenizer is already configured via the constructor's `tokenizer` parameter.
if not _VLLM_V0_12:
self._llm.set_tokenizer(self._tokenizer) # pyright: ignore[reportAttributeAccessIssue]
# Optional tool-call parser. Direct kwarg wins over model_params.
parser_name = tool_call_parser or model_params.tool_call_parser
self._tool_parser = None
if parser_name is not None:
if not _VLLM_TOOL_PARSERS_AVAILABLE:
raise RuntimeError(
"vLLM tool parsers are not available in this vLLM version. "
"Upgrade vLLM or unset `tool_call_parser`."
)
try:
parser_cls = ToolParserManager.get_tool_parser(parser_name) # pyright: ignore[reportOptionalMemberAccess]
except KeyError as e:
raise ValueError(
f"Unknown vLLM tool_call_parser '{parser_name}'."
) from e
self._tool_parser = parser_cls(self._tokenizer) # pyright: ignore[reportArgumentType]
logger.info(f"VLLM engine will parse tool calls with '{parser_name}'.")
@staticmethod
def _normalize_vllm_finish_reason(raw_reason: str | None) -> FinishReason | None:
"""Normalize vLLM finish_reason string to FinishReason enum."""
if raw_reason is None:
return None
mapping = {
"stop": FinishReason.STOP,
"length": FinishReason.LENGTH,
}
return mapping.get(raw_reason.lower(), FinishReason.UNKNOWN)
def _convert_conversation_to_vllm_input(
self, conversation: Conversation
) -> list[ChatCompletionMessageParam]:
"""Converts a conversation to a list of vllm input messages.
Args:
conversation: The conversation to convert.
Returns:
List[ChatCompletionMessageParam]: A list of vllm input messages.
"""
result: list[ChatCompletionMessageParam] = []
for json_dict in create_list_of_message_json_dicts(
conversation.messages, group_adjacent_same_role_turns=True
):
if "role" not in json_dict:
raise RuntimeError("The required field 'role' is missing!")
if "content" not in json_dict:
raise RuntimeError("The required field 'content' is missing!")
content = json_dict["content"]
# Assistant messages with only `tool_calls` legitimately have
# `content=None` per the OpenAI wire format.
if content is not None and not isinstance(content, str | list):
raise RuntimeError(
"The 'content' field must be `str`, `list`, or `None`. "
f"Actual: {type(content)}."
)
result.append(json_dict) # type: ignore[arg-type]
return result
def _infer(
self,
input: list[Conversation],
inference_config: InferenceConfig | None = None,
) -> list[Conversation]:
"""Runs model inference on the provided input.
Documentation: https://docs.vllm.ai/en/stable/dev/sampling_params.html
Args:
input: A list of conversations to run inference on.
inference_config: Parameters for inference.
Returns:
List[Conversation]: Inference output.
"""
generation_params = (
inference_config.generation
if inference_config and inference_config.generation
else self._generation_params
)
model_params = (
inference_config.model
if inference_config and inference_config.model
else self._model_params
)
if generation_params.guided_decoding is not None:
if _VLLM_V0_12:
# vLLM v0.12+ uses StructuredOutputsParams (direct construction)
guided_decoding = VLLMGuidedDecodingParams(
json=generation_params.guided_decoding.json,
regex=generation_params.guided_decoding.regex,
choice=generation_params.guided_decoding.choice,
)
else:
# vLLM <0.12 uses GuidedDecodingParams.from_optional()
guided_decoding = VLLMGuidedDecodingParams.from_optional( # type: ignore[attr-defined] # pyright: ignore[reportAttributeAccessIssue]
json=generation_params.guided_decoding.json,
regex=generation_params.guided_decoding.regex,
choice=generation_params.guided_decoding.choice,
)
else:
guided_decoding = None
# In vLLM v0.12+, the kwarg was renamed from 'guided_decoding'
# to 'structured_outputs'.
guided_decoding_kwarg = (
{"structured_outputs": guided_decoding}
if _VLLM_V0_12
else {"guided_decoding": guided_decoding}
)
sampling_params = SamplingParams(
n=1,
max_tokens=generation_params.max_new_tokens,
temperature=generation_params.temperature,
top_p=generation_params.top_p
if generation_params.top_p is not None
else 1.0,
frequency_penalty=generation_params.frequency_penalty,
presence_penalty=generation_params.presence_penalty,
stop=generation_params.stop_strings,
stop_token_ids=generation_params.stop_token_ids,
min_p=generation_params.min_p,
**guided_decoding_kwarg, # pyright: ignore[reportArgumentType]
skip_special_tokens=generation_params.skip_special_tokens,
)
output_conversations = []
vllm_conversations = []
non_skipped_conversations = []
for conversation in input:
if not conversation.messages:
logger.warning("Conversation must have at least one message.")
continue
vllm_input = self._convert_conversation_to_vllm_input(conversation)
vllm_conversations.append(vllm_input)
non_skipped_conversations.append(conversation)
if len(vllm_conversations) == 0:
return []
any_tools = any(c.tools for c in non_skipped_conversations)
if not any_tools:
# Fast path: single batched call, unchanged from prior behavior.
chat_responses = self._llm.chat(
vllm_conversations,
sampling_params=sampling_params,
lora_request=self._lora_request,
use_tqdm=(len(vllm_conversations) >= 2),
chat_template=None,
chat_template_content_format="auto",
chat_template_kwargs=model_params.chat_template_kwargs,
)
else:
# vLLM's `chat(tools=...)` applies one tool list to the whole
# batch, so we group conversations by identical tools and dispatch
# one chat() per group. Original input order is preserved.
chat_responses: list = [None] * len(non_skipped_conversations)
for indices, group_tools in self._group_by_tools(non_skipped_conversations):
group_inputs = [vllm_conversations[i] for i in indices]
group_responses = self._llm.chat(
group_inputs,
sampling_params=sampling_params,
lora_request=self._lora_request,
use_tqdm=(len(group_inputs) >= 2),
chat_template=None,
chat_template_content_format="auto",
chat_template_kwargs=model_params.chat_template_kwargs,
tools=group_tools,
)
for idx, resp in zip(indices, group_responses):
chat_responses[idx] = resp
for conversation, chat_response in zip(
non_skipped_conversations, chat_responses
):
assert chat_response is not None
new_messages, finish_reason_override = self._build_response_messages(
conversation, chat_response
)
messages = [
*conversation.messages,
*new_messages,
]
metadata = dict(conversation.metadata)
if chat_response.outputs:
if finish_reason_override is not None:
metadata["finish_reason"] = finish_reason_override.value
else:
raw_reason = chat_response.outputs[0].finish_reason
finish_reason = self._normalize_vllm_finish_reason(raw_reason)
if finish_reason is not None:
metadata["finish_reason"] = finish_reason.value
new_conversation = Conversation(
messages=messages,
metadata=metadata,
conversation_id=conversation.conversation_id,
)
self._save_conversation_to_scratch(
new_conversation,
inference_config.output_path if inference_config else None,
)
output_conversations.append(new_conversation)
return output_conversations
@staticmethod
def _group_by_tools(
conversations: list[Conversation],
) -> list[tuple[list[int], list[dict] | None]]:
"""Groups conversation indices by identical tool definitions.
Tools come from ``Conversation.tools`` (post-validation a list of
``ToolDefinition`` Pydantic models, or ``None``). vLLM's
``LLM.chat(tools=...)`` kwarg expects OpenAI-format dicts and is
batch-global, so we dump to dicts up front and split heterogeneous
batches.
Returns:
A list of ``(indices, tools)`` tuples in the order each unique
tools list is first seen.
"""
groups: list[tuple[list[int], list[dict] | None]] = []
key_to_group: dict[str, int] = {}
for i, conv in enumerate(conversations):
tools: list[dict] | None = (
None
if conv.tools is None
else [t.model_dump(mode="json", exclude_none=True) for t in conv.tools]
)
if tools is None:
key = "__none__"
else:
key = json.dumps(tools, sort_keys=True)
if key not in key_to_group:
key_to_group[key] = len(groups)
groups.append(([], tools))
groups[key_to_group[key]][0].append(i)
return groups
def _build_response_messages(
self,
conversation: Conversation,
chat_response,
) -> tuple[list[Message], FinishReason | None]:
"""Builds assistant messages from a vLLM chat response.
When ``self._tool_parser`` is set, runs it over each completion's
text and populates ``Message.tool_calls`` with the extracted payload.
On parser failure, falls back to raw text.
"""
new_messages: list[Message] = []
finish_reason_override: FinishReason | None = None
for completion in chat_response.outputs:
text = completion.text
content: str | None = text
tool_calls_payload: list[ToolCall] | None = None
if self._tool_parser is not None:
# Some parsers read `request.tool_choice` from the
# non-streaming entry point; pass a stub so they don't crash.
# vLLM offline `LLM.chat()` has no `tool_choice` knob, so
# "auto" is the only honest value — it matches OpenAI's
# default-when-unset and the offline reality of "model
# decides based on the rendered tools". Other values would
# mislead the parser ("none" would drop real calls;
# "required" or a specific function would force/filter
# interpretations the model didn't actually make).
stub = SimpleNamespace(
tool_choice="auto",
tools=[
t.model_dump(mode="json", exclude_none=True)
for t in (conversation.tools or [])
],
)
try:
extracted = self._tool_parser.extract_tool_calls(
text,
request=stub, # type: ignore[arg-type]
)
except Exception:
logger.exception(
"Tool-call parser %s failed; falling back to raw text.",
type(self._tool_parser).__name__,
)
extracted = None
if extracted is not None and getattr(extracted, "tools_called", False):
tool_calls_payload = [
ToolCall.model_validate(tc.model_dump())
for tc in extracted.tool_calls
]
# Empty leading content is rendered as `None` to match
# the OpenAI wire format for tool-only assistant turns.
content = extracted.content or None
finish_reason_override = FinishReason.TOOL_CALLS
new_messages.append(
Message(
content=content,
role=Role.ASSISTANT,
tool_calls=tool_calls_payload,
)
)
return new_messages, finish_reason_override
[docs]
def infer_online(
self,
input: list[Conversation],
inference_config: InferenceConfig | None = None,
) -> list[Conversation]:
"""Runs model inference online.
Args:
input: A list of conversations to run inference on.
inference_config: Parameters for inference.
Returns:
List[Conversation]: Inference output.
"""
warnings.warn(
"infer_online() will be private in the future. Use infer() instead.",
DeprecationWarning,
stacklevel=2,
)
results = self._infer_online(input, inference_config)
if inference_config and inference_config.output_path:
self._save_conversations(results, inference_config.output_path)
return results
[docs]
def infer_from_file(
self,
input_filepath: str,
inference_config: InferenceConfig | None = None,
) -> list[Conversation]:
"""Runs model inference on inputs in the provided file.
This is a convenience method to prevent boilerplate from asserting the
existence of input_filepath in the generation_params.
Args:
input_filepath: Path to the input file containing prompts for
generation.
inference_config: Parameters for inference.
Returns:
List[Conversation]: Inference output.
"""
warnings.warn(
"infer_from_file() will be private in the future. Use infer() instead.",
DeprecationWarning,
stacklevel=2,
)
input = self._read_conversations(input_filepath)
results = self._infer(input, inference_config)
if inference_config and inference_config.output_path:
self._save_conversations(results, inference_config.output_path)
return results
@override
def _infer_online(
self,
input: list[Conversation],
inference_config: InferenceConfig | None = None,
) -> list[Conversation]:
"""Runs model inference online.
Args:
input: A list of conversations to run inference on.
inference_config: Parameters for inference.
Returns:
List[Conversation]: Inference output.
"""
return self._infer(input, inference_config)
[docs]
@override
def get_supported_params(self) -> set[str]:
"""Returns a set of supported generation parameters for this engine."""
return {
"frequency_penalty",
"guided_decoding",
"max_new_tokens",
"min_p",
"presence_penalty",
"skip_special_tokens",
"stop_strings",
"stop_token_ids",
"temperature",
"top_p",
}
