43 lines
1.7 KiB
Python
43 lines
1.7 KiB
Python
from vllm.inputs.data import SingletonPrompt
|
|
from vllm.inputs.parse import (
|
|
ParsedEmbedsPrompt,
|
|
ParsedSingletonPrompt,
|
|
ParsedStrPrompt,
|
|
ParsedTextPrompt,
|
|
ParsedTokensPrompt,
|
|
)
|
|
|
|
|
|
def parse_singleton_prompt_omni(prompt: SingletonPrompt) -> ParsedSingletonPrompt:
|
|
"""Parse a singleton prompt into a typed parsed prompt.
|
|
|
|
Handles omni-specific prompt types including tokens prompts with
|
|
embeddings and additional information. Supports string, text,
|
|
tokens, and embeddings prompts.
|
|
|
|
Args:
|
|
prompt: Singleton prompt to parse. Can be a string, TextPrompt,
|
|
TokensPrompt (with optional prompt_embeds and additional_information),
|
|
or EmbedsPrompt.
|
|
|
|
Returns:
|
|
ParsedSingletonPrompt containing the parsed prompt with type information
|
|
|
|
Raises:
|
|
TypeError: If the prompt type is not supported
|
|
"""
|
|
if isinstance(prompt, str):
|
|
return ParsedStrPrompt(type="str", content=prompt)
|
|
elif isinstance(prompt, dict):
|
|
# Type ignores are because mypy does not correctly infer the TypedDicts
|
|
# Pyright does succeed.
|
|
# Priority tokens: When both tokens and embeds exist, keep both and
|
|
# follow the tokens path
|
|
if "prompt_token_ids" in prompt:
|
|
return ParsedTokensPrompt(type="tokens", content=prompt) # type: ignore[typeddict-item]
|
|
elif "prompt_embeds" in prompt:
|
|
return ParsedEmbedsPrompt(type="embeds", content=prompt) # type: ignore[typeddict-item]
|
|
elif "prompt" in prompt:
|
|
return ParsedTextPrompt(type="text", content=prompt)
|
|
raise TypeError("inputs must be a string, TextPrompt, TokensPrompt, or EmbedsPrompt")
|