Spaces:
Running
Running
| from .base_handler import ModelHandler | |
| from .nlp_models.sequence_classification_handler import SequenceClassificationHandler | |
| from .nlp_models.question_answering_handler import QuestionAnsweringHandler | |
| from .nlp_models.token_classification_handler import TokenClassificationHandler | |
| from .nlp_models.causal_lm_handler import CausalLMHandler | |
| from .nlp_models.embedding_model_handler import EmbeddingModelHandler | |
| from .audio_models.whisper_handler import WhisperHandler | |
| from .nlp_models.masked_lm_handler import MaskedLMHandler | |
| from .nlp_models.seq2seq_lm_handler import Seq2SeqLMHandler | |
| from .nlp_models.multiple_choice_handler import MultipleChoiceHandler | |
| from .img_models.image_classification_handler import ImageClassificationHandler | |
| from transformers import ( | |
| AutoModel, | |
| AutoModelForTokenClassification, | |
| AutoModelForSequenceClassification, | |
| AutoModelForQuestionAnswering, | |
| AutoModelForCausalLM, | |
| AutoModelForMaskedLM, | |
| AutoModelForSeq2SeqLM, | |
| AutoModelForMultipleChoice, | |
| ) | |
| TASK_CONFIGS = { | |
| "embedding": { | |
| "model_class": AutoModel, | |
| "handler_class": EmbeddingModelHandler, | |
| "example_text": "Hey, I am feeling way to good to be true.", | |
| }, | |
| "ner": { | |
| "model_class": AutoModelForTokenClassification, | |
| "handler_class": TokenClassificationHandler, | |
| "example_text": "John works at Google in New York as a software engineer.", | |
| }, | |
| "text_classification": { | |
| "model_class": AutoModelForSequenceClassification, | |
| "handler_class": SequenceClassificationHandler, | |
| "example_text": "This movie was great and I loved it.", | |
| }, | |
| "question_answering": { | |
| "model_class": AutoModelForQuestionAnswering, | |
| "handler_class": QuestionAnsweringHandler, | |
| "example_text": "The pyramids were built in ancient Egypt. QUES: Where were the pyramids built?", | |
| }, | |
| "causal_lm": { | |
| "model_class": AutoModelForCausalLM, | |
| "handler_class": CausalLMHandler, | |
| "example_text": "Once upon a time, there was ", | |
| }, | |
| "mask_lm": { | |
| "model_class": AutoModelForMaskedLM, | |
| "handler_class": MaskedLMHandler, | |
| "example_text": "The quick brown [MASK] jumps over the lazy dog.", | |
| }, | |
| "seq2seq_lm": { | |
| "model_class": AutoModelForSeq2SeqLM, | |
| "handler_class": Seq2SeqLMHandler, | |
| "example_text": "Translate English to French: The house is wonderful.", | |
| }, | |
| "multiple_choice": { | |
| "model_class": AutoModelForMultipleChoice, | |
| "handler_class": MultipleChoiceHandler, | |
| "example_text": "What is the capital of France? (A) Paris (B) London (C) Berlin (D) Rome", | |
| }, | |
| "whisper_finetuning": { | |
| "model_class": None, # Not implemented | |
| "handler_class": WhisperHandler, | |
| "example_text": "!!!!!NOT IMPLEMENTED!!!!!", | |
| }, | |
| "image_classification": { | |
| "model_class": None, # Not implemented | |
| "handler_class": ImageClassificationHandler, | |
| "example_text": "!!!!!NOT IMPLEMENTED!!!!!", | |
| }, | |
| } | |
| def get_model_handler(task: str, model_name: str, quantization_type: str, test_text: str): | |
| task_config = TASK_CONFIGS.get(task) | |
| if not task_config: | |
| raise ValueError(f"No configuration found for task: {task}") | |
| handler_class = task_config["handler_class"] | |
| model_class = task_config["model_class"] | |
| return handler_class(model_name, model_class, quantization_type, test_text) |