# Autogenerated code do not modify manually.
# Run `poetry run python generate_enums.py` to update.
from enum import Enum, StrEnum
from typing import Any, Self
class ModelsEnum(Enum):
@classmethod
def get(cls, key: str | Self) -> Self:
if isinstance(key, ModelsEnum):
return key
name = key.upper().replace(" ", "_").replace("_-_", "_").replace("-", "_")
return cls[name]
@classmethod
def values(cls) -> list[Any]:
return [x.value for x in cls]
@classmethod
def keys(cls) -> list[str]:
return [x.name for x in cls]
class ModelsStrEnum(StrEnum, ModelsEnum):
pass
[docs]
class ArtifactType(ModelsStrEnum):
CHARIOT = "chariot"
HUGGINGFACE = "huggingface"
NEURALMAGIC = "neuralmagic"
ONNX = "onnx"
PYTORCH = "pytorch"
SKLEARN = "sklearn"
[docs]
class ArtifactTypesTaskType(ModelsEnum):
CHARIOT = ["Image Classification", "Image Embedding", "Image Segmentation", "Object Detection"]
HUGGINGFACE = [
"Automatic Speech Recognition",
"Conversational",
"Summarization",
"Text Classification",
"Text Embedding",
"Text Fill-Mask",
"Text Generation",
"Token Classification",
"Translation",
]
NEURALMAGIC = ["Image Classification"]
ONNX = [
"Automatic Speech Recognition",
"Conversational",
"Feature Extraction",
"Image Autoencoder",
"Image Classification",
"Image Embedding",
"Image Generation",
"Image Segmentation",
"Object Detection",
"Oriented Object Detection",
"Other - Computer Vision",
"Other - Natural Language",
"Other - Structured Data",
"Question Answer",
"Structured Data Classification",
"Structured Data Regression",
"Summarization",
"Text Classification",
"Text Embedding",
"Text Fill-Mask",
"Text Generation",
"Text2Text Generation",
"Token Classification",
"Translation",
]
PYTORCH = [
"Automatic Speech Recognition",
"Conversational",
"Feature Extraction",
"Image Autoencoder",
"Image Classification",
"Image Embedding",
"Image Generation",
"Image Segmentation",
"Object Detection",
"Oriented Object Detection",
"Other - Computer Vision",
"Other - Natural Language",
"Other - Structured Data",
"Question Answer",
"Structured Data Classification",
"Structured Data Regression",
"Summarization",
"Text Classification",
"Text Embedding",
"Text Fill-Mask",
"Text Generation",
"Text2Text Generation",
"Token Classification",
"Translation",
]
SKLEARN = [
"Other - Structured Data",
"Structured Data Classification",
"Structured Data Regression",
]
[docs]
class InferenceEngine(ModelsStrEnum):
CHARIOTDEEPSPARSE = "ChariotDeepSparse"
CHARIOTPYTORCH = "ChariotPytorch"
HUGGINGFACE = "Huggingface"
VLLM = "vLLM"
[docs]
class Protocol(ModelsStrEnum):
V2 = "v2"
[docs]
class TaskType(ModelsStrEnum):
AUTOMATIC_SPEECH_RECOGNITION = "Automatic Speech Recognition"
CONVERSATIONAL = "Conversational"
FEATURE_EXTRACTION = "Feature Extraction"
IMAGE_AUTOENCODER = "Image Autoencoder"
IMAGE_CLASSIFICATION = "Image Classification"
IMAGE_EMBEDDING = "Image Embedding"
IMAGE_GENERATION = "Image Generation"
IMAGE_SEGMENTATION = "Image Segmentation"
OBJECT_DETECTION = "Object Detection"
ORIENTED_OBJECT_DETECTION = "Oriented Object Detection"
OTHER_COMPUTER_VISION = "Other - Computer Vision"
OTHER_NATURAL_LANGUAGE = "Other - Natural Language"
OTHER_STRUCTURED_DATA = "Other - Structured Data"
QUESTION_ANSWER = "Question Answer"
STRUCTURED_DATA_CLASSIFICATION = "Structured Data Classification"
STRUCTURED_DATA_REGRESSION = "Structured Data Regression"
SUMMARIZATION = "Summarization"
TEXT_CLASSIFICATION = "Text Classification"
TEXT_EMBEDDING = "Text Embedding"
TEXT_FILL_MASK = "Text Fill-Mask"
TEXT_GENERATION = "Text Generation"
TEXT2TEXT_GENERATION = "Text2Text Generation"
TOKEN_CLASSIFICATION = "Token Classification"
TRANSLATION = "Translation"
[docs]
class TaskTypesInferenceMethod(ModelsEnum):
AUTOMATIC_SPEECH_RECOGNITION = ["predict"]
CONVERSATIONAL = ["chat"]
FEATURE_EXTRACTION = ["predict"]
IMAGE_AUTOENCODER = ["embed", "reconstruct"]
IMAGE_CLASSIFICATION = ["embed", "predict", "predict_proba"]
IMAGE_EMBEDDING = ["embed"]
IMAGE_GENERATION = ["predict"]
IMAGE_SEGMENTATION = ["predict", "predict_proba"]
OBJECT_DETECTION = ["detect"]
ORIENTED_OBJECT_DETECTION = ["detect"]
OTHER_COMPUTER_VISION = ["predict"]
OTHER_NATURAL_LANGUAGE = ["predict"]
OTHER_STRUCTURED_DATA = ["predict"]
QUESTION_ANSWER = ["predict"]
STRUCTURED_DATA_CLASSIFICATION = ["predict"]
STRUCTURED_DATA_REGRESSION = ["predict"]
SUMMARIZATION = ["predict"]
TEXT_CLASSIFICATION = ["predict"]
TEXT_EMBEDDING = ["embed"]
TEXT_FILL_MASK = ["predict"]
TEXT_GENERATION = ["complete"]
TEXT2TEXT_GENERATION = ["predict"]
TOKEN_CLASSIFICATION = ["predict"]
TRANSLATION = ["predict"]
[docs]
class TaskTypesRequirement(ModelsEnum):
AUTOMATIC_SPEECH_RECOGNITION = {
"class_labels": False,
"input_info": False,
"input_modality": "text",
}
CONVERSATIONAL = {"class_labels": False, "input_info": False, "input_modality": "text"}
FEATURE_EXTRACTION = {"class_labels": False, "input_info": False, "input_modality": "text"}
IMAGE_AUTOENCODER = {"class_labels": False, "input_info": False, "input_modality": "image"}
IMAGE_CLASSIFICATION = {"class_labels": True, "input_info": False, "input_modality": "image"}
IMAGE_EMBEDDING = {"class_labels": False, "input_info": False, "input_modality": "image"}
IMAGE_GENERATION = {"class_labels": False, "input_info": False, "input_modality": "image"}
IMAGE_SEGMENTATION = {"class_labels": True, "input_info": False, "input_modality": "image"}
OBJECT_DETECTION = {"class_labels": True, "input_info": False, "input_modality": "image"}
ORIENTED_OBJECT_DETECTION = {
"class_labels": True,
"input_info": False,
"input_modality": "image",
}
OTHER_COMPUTER_VISION = {"class_labels": False, "input_info": False, "input_modality": "image"}
OTHER_NATURAL_LANGUAGE = {"class_labels": False, "input_info": False, "input_modality": "text"}
OTHER_STRUCTURED_DATA = {"class_labels": False, "input_info": True, "input_modality": "tabular"}
QUESTION_ANSWER = {"class_labels": False, "input_info": False, "input_modality": "text"}
STRUCTURED_DATA_CLASSIFICATION = {
"class_labels": True,
"input_info": True,
"input_modality": "tabular",
}
STRUCTURED_DATA_REGRESSION = {
"class_labels": False,
"input_info": True,
"input_modality": "tabular",
}
SUMMARIZATION = {"class_labels": False, "input_info": False, "input_modality": "text"}
TEXT_CLASSIFICATION = {"class_labels": True, "input_info": False, "input_modality": "text"}
TEXT_EMBEDDING = {"class_labels": False, "input_info": False, "input_modality": "text"}
TEXT_FILL_MASK = {"class_labels": False, "input_info": False, "input_modality": "text"}
TEXT_GENERATION = {"class_labels": False, "input_info": False, "input_modality": "text"}
TEXT2TEXT_GENERATION = {"class_labels": False, "input_info": False, "input_modality": "text"}
TOKEN_CLASSIFICATION = {"class_labels": True, "input_info": False, "input_modality": "text"}
TRANSLATION = {"class_labels": False, "input_info": False, "input_modality": "text"}
__all__ = [
"ArtifactType",
"ArtifactTypesTaskType",
"InferenceEngine",
"Protocol",
"TaskType",
"TaskTypesInferenceMethod",
"TaskTypesRequirement",
]