from .imports import *
from .constants import *
# pipeline_tag -> task vocabulary. Every value is HF-derivable.
HF_TASK_TO_TASKS = {
    "text-generation": ["text-generation"],
    "image-text-to-text": ["image-text-to-text", "text-generation"],
    "automatic-speech-recognition": ["automatic-speech-recognition"],
    "text-summarization": ["text-summarization"],
    "text2text-generation": ["text-summarization", "text2text-generation"],
    # KeyBERT rides any sentence-transformers model, so embedding models
    # also serve keyword-extraction.
    "feature-extraction": ["feature-extraction", "sentence-similarity", "keyword-extraction"],
    "sentence-similarity": ["feature-extraction", "sentence-similarity", "keyword-extraction"],
    "text-to-image": ["text-to-image"],
}
RUNNER_PAIRS = {
    ("transformers", "text-generation"), ("llama_cpp", "text-generation"),
    ("transformers", "image-text-to-text"), ("llama_cpp", "image-text-to-text"),
    ("transformers", "automatic-speech-recognition"),
    ("transformers", "text-summarization"), ("transformers", "text2text-generation"),
    ("transformers", "feature-extraction"), ("transformers", "sentence-similarity"),
    ("transformers", "text-to-image"), ("transformers", "keyword-extraction"),
}
MEDIA_DEFAULTS: Dict[str, str] = {
    "document": DEFAULT_CHAT_MODEL,
    "code":     DEFAULT_CHAT_MODEL,
    "text":     DEFAULT_CHAT_MODEL,
    "image":    DEFAULT_VISION_MODEL,
    "audio":    DEFAULT_WHISPER_MODEL,
    "video":    DEFAULT_WHISPER_MODEL,
}

TASK_DEFAULTS: Dict[str, str] = {
    "text-generation":              DEFAULT_CHAT_MODEL,
    "image-text-to-text":           DEFAULT_VISION_MODEL,
    "automatic-speech-recognition": DEFAULT_WHISPER_MODEL,
    "text-summarization":           DEFAULT_SUMMARIZE_MODEL,
    "text2text-generation":         DEFAULT_SUMMARIZE_MODEL,
    "feature-extraction":           DEFAULT_EMBED_MODEL,
    "sentence-similarity":          DEFAULT_EMBED_MODEL,
    "text-to-image":                DEFAULT_IMAGEGEN_MODEL,
    "keyword-extraction":           DEFAULT_KEYWORDS_MODEL,
}
