68 lines
2.1 KiB
Python
68 lines
2.1 KiB
Python
from typing import Literal, TypedDict
|
|
|
|
from vllm.inputs import EmbedsPrompt, SingletonPrompt, TextPrompt, TokensPrompt
|
|
|
|
|
|
class ParsedStrPrompt(TypedDict):
|
|
type: Literal["str"]
|
|
content: str
|
|
|
|
|
|
class ParsedTextPrompt(TypedDict):
|
|
type: Literal["text"]
|
|
content: TextPrompt
|
|
|
|
|
|
class ParsedTokensPrompt(TypedDict):
|
|
type: Literal["tokens"]
|
|
content: TokensPrompt
|
|
|
|
|
|
class ParsedEmbedsPrompt(TypedDict):
|
|
type: Literal["embeds"]
|
|
content: EmbedsPrompt
|
|
|
|
|
|
ParsedSingletonPrompt = ParsedStrPrompt | ParsedTextPrompt | ParsedTokensPrompt | ParsedEmbedsPrompt
|
|
|
|
|
|
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 {"type": "str", "content": prompt}
|
|
if isinstance(prompt, list):
|
|
if not all(isinstance(token, int) for token in prompt):
|
|
raise TypeError("Token prompt should be a list of integers")
|
|
return {
|
|
"type": "tokens",
|
|
"content": TokensPrompt(prompt_token_ids=prompt),
|
|
}
|
|
if isinstance(prompt, dict):
|
|
# Priority tokens: When both tokens and embeds exist, keep both and
|
|
# follow the tokens path.
|
|
if "prompt_token_ids" in prompt:
|
|
return {"type": "tokens", "content": prompt}
|
|
if "prompt_embeds" in prompt:
|
|
return {"type": "embeds", "content": prompt}
|
|
if "prompt" in prompt:
|
|
return {"type": "text", "content": prompt}
|
|
raise TypeError("inputs must be a string, list of token IDs, TextPrompt, TokensPrompt, or EmbedsPrompt")
|