opsml.types.huggingface

 1from enum import Enum
 2
 3
 4class HuggingFaceTask(str, Enum):
 5    AUDIO_CLASSIFICATION = "audio-classification"
 6    AUTOMATIC_SPEECH_RECOGNITION = "automatic-speech-recognition"
 7    CONVERSATIONAL = "conversational"
 8    DEPTH_ESTIMATION = "depth-estimation"
 9    DOCUMENT_QUESTION_ANSWERING = "document-question-answering"
10    FEATURE_EXTRACTION = "feature-extraction"
11    FILL_MASK = "fill-mask"
12    IMAGE_CLASSIFICATION = "image-classification"
13    IMAGE_SEGMENTATION = "image-segmentation"
14    IMAGE_TO_IMAGE = "image-to-image"
15    IMAGE_TO_TEXT = "image-to-text"
16    MASK_GENERATION = "mask-generation"
17    OBJECT_DETECTION = "object-detection"
18    QUESTION_ANSWERING = "question-answering"
19    SUMMARIZATION = "summarization"
20    TABLE_QUESTION_ANSWERING = "table-question-answering"
21    TEXT2TEXT_GENERATION = "text2text-generation"
22    TEXT_CLASSIFICATION = "text-classification"
23    TEXT_GENERATION = "text-generation"
24    TEXT_TO_AUDIO = "text-to-audio"
25    TOKEN_CLASSIFICATION = "token-classification"
26    TRANSLATION = "translation"
27    TRANSLATION_XX_TO_YY = "translation_xx_to_yy"
28    VIDEO_CLASSIFICATION = "video-classification"
29    VISUAL_QUESTION_ANSWERING = "visual-question-answering"
30    ZERO_SHOT_CLASSIFICATION = "zero-shot-classification"
31    ZERO_SHOT_IMAGE_CLASSIFICATION = "zero-shot-image-classification"
32    ZERO_SHOT_AUDIO_CLASSIFICATION = "zero-shot-audio-classification"
33    ZERO_SHOT_OBJECT_DETECTION = "zero-shot-object-detection"
34
35
36GENERATION_TYPES = [
37    HuggingFaceTask.MASK_GENERATION.value,
38    HuggingFaceTask.TEXT_GENERATION.value,
39    HuggingFaceTask.TEXT2TEXT_GENERATION.value,
40]
41
42
43class HuggingFaceORTModel(str, Enum):
44    ORT_AUDIO_CLASSIFICATION = "ORTModelForAudioClassification"
45    ORT_AUDIO_FRAME_CLASSIFICATION = "ORTModelForAudioFrameClassification"
46    ORT_AUDIO_XVECTOR = "ORTModelForAudioXVector"
47    ORT_CUSTOM_TASKS = "ORTModelForCustomTasks"
48    ORT_CTC = "ORTModelForCTC"
49    ORT_FEATURE_EXTRACTION = "ORTModelForFeatureExtraction"
50    ORT_IMAGE_CLASSIFICATION = "ORTModelForImageClassification"
51    ORT_MASKED_LM = "ORTModelForMaskedLM"
52    ORT_MULTIPLE_CHOICE = "ORTModelForMultipleChoice"
53    ORT_QUESTION_ANSWERING = "ORTModelForQuestionAnswering"
54    ORT_SEMANTIC_SEGMENTATION = "ORTModelForSemanticSegmentation"
55    ORT_SEQUENCE_CLASSIFICATION = "ORTModelForSequenceClassification"
56    ORT_TOKEN_CLASSIFICATION = "ORTModelForTokenClassification"
57    ORT_SEQ2SEQ_LM = "ORTModelForSeq2SeqLM"
58    ORT_SPEECH_SEQ2SEQ = "ORTModelForSpeechSeq2Seq"
59    ORT_VISION2SEQ = "ORTModelForVision2Seq"
60    ORT_PIX2STRUCT = "ORTModelForPix2Struct"
61    ORT_CAUSAL_LM = "ORTModelForCausalLM"
62    ORT_OPTIMIZER = "ORTOptimizer"
63    ORT_QUANTIZER = "ORTQuantizer"
64    ORT_TRAINER = "ORTTrainer"
65    ORT_SEQ2SEQ_TRAINER = "ORTSeq2SeqTrainer"
66    ORT_TRAINING_ARGUMENTS = "ORTTrainingArguments"
67    ORT_SEQ2SEQ_TRAINING_ARGUMENTS = "ORTSeq2SeqTrainingArguments"
68    ORT_STABLE_DIFFUSION_PIPELINE = "ORTStableDiffusionPipeline"
69    ORT_STABLE_DIFFUSION_IMG2IMG_PIPELINE = "ORTStableDiffusionImg2ImgPipeline"
70    ORT_STABLE_DIFFUSION_INPAINT_PIPELINE = "ORTStableDiffusionInpaintPipeline"
71    ORT_STABLE_DIFFUSION_XL_PIPELINE = "ORTStableDiffusionXLPipeline"
72    ORT_STABLE_DIFFUSION_XL_IMG2IMG_PIPELINE = "ORTStableDiffusionXLImg2ImgPipeline"