32 lines
1.1 KiB
Python
32 lines
1.1 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:
|
|
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"
|
|
)
|