Module: Informers
- Defined in:
- lib/informers/utils/ffmpeg.rb,
lib/informers.rb,
lib/informers/env.rb,
lib/informers/model.rb,
lib/informers/models.rb,
lib/informers/configs.rb,
lib/informers/version.rb,
lib/informers/pipelines.rb,
lib/informers/utils/hub.rb,
lib/informers/processors.rb,
lib/informers/tokenizers.rb,
lib/informers/utils/core.rb,
lib/informers/utils/math.rb,
lib/informers/utils/audio.rb,
lib/informers/utils/image.rb,
lib/informers/utils/tensor.rb,
lib/informers/utils/generation.rb
Overview
Copyright 2021 The HuggingFace Team. All rights reserved.
Licensed under the Apache License, Version 2.0 (the “License”); you may not use this file except in compliance with the License. You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an “AS IS” BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.
Defined Under Namespace
Modules: Utils Classes: AudioClassificationPipeline, AutoConfig, AutoModel, AutoModelForAudioClassification, AutoModelForAudioFrameClassification, AutoModelForCTC, AutoModelForCausalLM, AutoModelForDepthEstimation, AutoModelForDocumentQuestionAnswering, AutoModelForImageClassification, AutoModelForImageFeatureExtraction, AutoModelForImageMatting, AutoModelForImageSegmentation, AutoModelForImageToImage, AutoModelForMaskGeneration, AutoModelForMaskedLM, AutoModelForObjectDetection, AutoModelForQuestionAnswering, AutoModelForSemanticSegmentation, AutoModelForSeq2SeqLM, AutoModelForSequenceClassification, AutoModelForSpeechSeq2Seq, AutoModelForTextToSpectrogram, AutoModelForTextToWaveform, AutoModelForTokenClassification, AutoModelForVision2Seq, AutoModelForXVector, AutoModelForZeroShotObjectDetection, AutoProcessor, AutoTokenizer, AutomaticSpeechRecognitionPipeline, BartForConditionalGeneration, BartForSequenceClassification, BartModel, BartPretrainedModel, BartTokenizer, BertForMaskedLM, BertForSequenceClassification, BertForTokenClassification, BertModel, BertPreTrainedModel, BertTokenizer, CLIPFeatureExtractor, CLIPModel, CLIPPreTrainedModel, CLIPTokenizer, ClapFeatureExtractor, ClapModel, ClapPreTrainedModel, DPTFeatureExtractor, DPTForDepthEstimation, DPTModel, DPTPreTrainedModel, DebertaV2Model, DebertaV2PreTrainedModel, DebertaV2Tokenizer, DepthEstimationPipeline, DetrFeatureExtractor, DetrForObjectDetection, DetrForSegmentation, DetrModel, DetrObjectDetectionOutput, DetrPreTrainedModel, DetrSegmentationOutput, DistilBertForQuestionAnswering, DistilBertForSequenceClassification, DistilBertModel, DistilBertPreTrainedModel, DistilBertTokenizer, DocumentQuestionAnsweringPipeline, DonutFeatureExtractor, DonutSwinModel, DonutSwinPreTrainedModel, EmbeddingPipeline, Error, FeatureExtractionPipeline, FeatureExtractor, FillMaskPipeline, GPT2LMHeadModel, GPT2Model, GPT2PreTrainedModel, GPT2Tokenizer, ImageClassificationPipeline, ImageFeatureExtractionPipeline, ImageFeatureExtractor, ImageSegmentationPipeline, ImageToImagePipeline, ImageToTextPipeline, M2M100ForConditionalGeneration, M2M100Model, M2M100PreTrainedModel, M2M100Tokenizer, MBartForCausalLM, MBartModel, MBartPreTrainedModel, MPNetModel, MPNetPreTrainedModel, MPNetTokenizer, MaskedLMOutput, Model, ModelOutput, NllbTokenizer, NomicBertModel, NomicBertPreTrainedModel, ObjectDetectionPipeline, OwlViTFeatureExtractor, OwlViTForObjectDetection, OwlViTModel, OwlViTPreTrainedModel, Pipeline, PreTrainedModel, PreTrainedTokenizer, PretrainedConfig, PretrainedMixin, Processor, QuestionAnsweringModelOutput, QuestionAnsweringPipeline, RerankingPipeline, RobertaForMaskedLM, RobertaModel, RobertaPreTrainedModel, RobertaTokenizer, Seq2SeqLMOutput, SequenceClassifierOutput, SpeechT5ForSpeechToText, SpeechT5ForTextToSpeech, SpeechT5Model, SpeechT5PreTrainedModel, SpeechT5Tokenizer, SummarizationPipeline, Swin2SRForImageSuperResolution, Swin2SRImageProcessor, Swin2SRModel, Swin2SRPreTrainedModel, T5ForConditionalGeneration, T5Model, T5PreTrainedModel, T5Tokenizer, Text2TextGenerationPipeline, TextClassificationPipeline, TextGenerationPipeline, TextToAudioPipeline, Todo, TokenClassificationPipeline, TokenClassifierOutput, TranslationPipeline, ViTFeatureExtractor, ViTForImageClassification, ViTModel, ViTPreTrainedModel, VisionEncoderDecoderModel, VitsModel, VitsPreTrainedModel, Wav2Vec2FeatureExtractor, Wav2Vec2ForSequenceClassification, Wav2Vec2Model, Wav2Vec2PreTrainedModel, WhisperFeatureExtractor, WhisperForConditionalGeneration, WhisperModel, WhisperPreTrainedModel, XLMRobertaForSequenceClassification, XLMRobertaModel, XLMRobertaPreTrainedModel, XLMRobertaTokenizer, ZeroShotAudioClassificationPipeline, ZeroShotClassificationPipeline, ZeroShotImageClassificationPipeline, ZeroShotObjectDetectionPipeline
Constant Summary collapse
- CACHE_HOME =
ENV.fetch("XDG_CACHE_HOME", File.join(ENV.fetch("HOME"), ".cache"))
- DEFAULT_CACHE_DIR =
File.(File.join(CACHE_HOME, "informers"))
- MODEL_TYPES =
{ EncoderOnly: 0, EncoderDecoder: 1, Seq2Seq: 2, Vision2Seq: 3, DecoderOnly: 4, MaskGeneration: 5 }
- MODEL_TYPE_MAPPING =
NOTE: These will be populated fully later
{}
- MODEL_NAME_TO_CLASS_MAPPING =
{}
- MODEL_CLASS_TO_NAME_MAPPING =
{}
- MODEL_MAPPING_NAMES_ENCODER_ONLY =
{ "bert" => ["BertModel", BertModel], "nomic_bert" => ["NomicBertModel", NomicBertModel], "deberta-v2" => ["DebertaV2Model", DebertaV2Model], "mpnet" => ["MPNetModel", MPNetModel], "distilbert" => ["DistilBertModel", DistilBertModel], "roberta" => ["RobertaModel", RobertaModel], "xlm-roberta" => ["XLMRobertaModel", XLMRobertaModel], "clap" => ["ClapModel", ClapModel], "clip" => ["CLIPModel", CLIPModel], "detr" => ["DetrModel", DetrModel], "vit" => ["ViTModel", ViTModel], "owlvit" => ["OwlViTModel", OwlViTModel], "donut-swin" => ["DonutSwinModel", DonutSwinModel] }
- MODEL_MAPPING_NAMES_ENCODER_DECODER =
{ "bart" => ["BartModel", BartModel] }
- MODEL_MAPPING_NAMES_DECODER_ONLY =
{ }
- MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES =
{ "whisper" => ["WhisperForConditionalGeneration", WhisperForConditionalGeneration] }
- MODEL_FOR_TEXT_TO_SPECTROGRAM_MAPPING_NAMES =
{ "speecht5" => ["SpeechT5ForTextToSpeech", SpeechT5ForTextToSpeech] }
- MODEL_FOR_TEXT_TO_WAVEFORM_MAPPING_NAMES =
{ "vits" => ["VitsModel", VitsModel] }
- MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES =
{ "bert" => ["BertForSequenceClassification", BertForSequenceClassification], "distilbert" => ["DistilBertForSequenceClassification", DistilBertForSequenceClassification], "xlm-roberta" => ["XLMRobertaForSequenceClassification", XLMRobertaForSequenceClassification], "bart" => ["BartForSequenceClassification", BartForSequenceClassification] }
- MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES =
{ "bert" => ["BertForTokenClassification", BertForTokenClassification] }
- MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES =
{ "t5" => ["T5ForConditionalGeneration", T5ForConditionalGeneration], "bart" => ["BartForConditionalGeneration", BartForConditionalGeneration], "m2m_100" => ["M2M100ForConditionalGeneration", M2M100ForConditionalGeneration] }
- MODEL_WITH_LM_HEAD_MAPPING_NAMES =
{ "gpt2" => ["GPT2LMHeadModel", GPT2LMHeadModel], "mbart" => ["MBartForCausalLM", MBartForCausalLM] }
- MODEL_FOR_MASKED_LM_MAPPING_NAMES =
{ "bert" => ["BertForMaskedLM", BertForMaskedLM], "roberta" => ["RobertaForMaskedLM", RobertaForMaskedLM] }
- MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES =
{ "distilbert" => ["DistilBertForQuestionAnswering", DistilBertForQuestionAnswering] }
- MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES =
{ "vision-encoder-decoder" => ["VisionEncoderDecoderModel", VisionEncoderDecoderModel] }
- MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING_NAMES =
{ "vision-encoder-decoder" => ["VisionEncoderDecoderModel", VisionEncoderDecoderModel] }
- MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES =
{ "vit" => ["ViTForImageClassification", ViTForImageClassification] }
- MODEL_FOR_OBJECT_DETECTION_MAPPING_NAMES =
{ "detr" => ["DetrForObjectDetection", DetrForObjectDetection] }
- MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING_NAMES =
{ "owlvit" => ["OwlViTForObjectDetection", OwlViTForObjectDetection] }
- MODEL_FOR_IMAGE_SEGMENTATION_MAPPING_NAMES =
{ "detr" => ["DetrForSegmentation", DetrForSegmentation] }
- MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING_NAMES =
{ }
- MODEL_FOR_MASK_GENERATION_MAPPING_NAMES =
{ }
- MODEL_FOR_CTC_MAPPING_NAMES =
{ }
- MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES =
{ "wav2vec2" => ["Wav2Vec2ForSequenceClassification", Wav2Vec2ForSequenceClassification] }
- MODEL_FOR_AUDIO_XVECTOR_MAPPING_NAMES =
{ }
- MODEL_FOR_AUDIO_FRAME_CLASSIFICATION_MAPPING_NAMES =
{ }
- MODEL_FOR_IMAGE_MATTING_MAPPING_NAMES =
{ }
- MODEL_FOR_IMAGE_TO_IMAGE_MAPPING_NAMES =
{ "swin2sr" => ["Swin2SRForImageSuperResolution", Swin2SRForImageSuperResolution] }
- MODEL_FOR_DEPTH_ESTIMATION_MAPPING_NAMES =
{ "dpt" => ["DPTForDepthEstimation", DPTForDepthEstimation] }
- MODEL_FOR_IMAGE_FEATURE_EXTRACTION_MAPPING_NAMES =
{ }
- MODEL_CLASS_TYPE_MAPPING =
[ [MODEL_MAPPING_NAMES_ENCODER_ONLY, MODEL_TYPES[:EncoderOnly]], [MODEL_MAPPING_NAMES_ENCODER_DECODER, MODEL_TYPES[:EncoderDecoder]], [MODEL_MAPPING_NAMES_DECODER_ONLY, MODEL_TYPES[:DecoderOnly]], [MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES, MODEL_TYPES[:EncoderOnly]], [MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES, MODEL_TYPES[:EncoderOnly]], [MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES, MODEL_TYPES[:Seq2Seq]], [MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES, MODEL_TYPES[:Seq2Seq]], [MODEL_WITH_LM_HEAD_MAPPING_NAMES, MODEL_TYPES[:DecoderOnly]], [MODEL_FOR_MASKED_LM_MAPPING_NAMES, MODEL_TYPES[:EncoderOnly]], [MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES, MODEL_TYPES[:EncoderOnly]], [MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES, MODEL_TYPES[:Vision2Seq]], [MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES, MODEL_TYPES[:EncoderOnly]], [MODEL_FOR_IMAGE_SEGMENTATION_MAPPING_NAMES, MODEL_TYPES[:EncoderOnly]], [MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING_NAMES, MODEL_TYPES[:EncoderOnly]], [MODEL_FOR_IMAGE_MATTING_MAPPING_NAMES, MODEL_TYPES[:EncoderOnly]], [MODEL_FOR_IMAGE_TO_IMAGE_MAPPING_NAMES, MODEL_TYPES[:EncoderOnly]], [MODEL_FOR_DEPTH_ESTIMATION_MAPPING_NAMES, MODEL_TYPES[:EncoderOnly]], [MODEL_FOR_OBJECT_DETECTION_MAPPING_NAMES, MODEL_TYPES[:EncoderOnly]], [MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING_NAMES, MODEL_TYPES[:EncoderOnly]], [MODEL_FOR_MASK_GENERATION_MAPPING_NAMES, MODEL_TYPES[:MaskGeneration]], [MODEL_FOR_CTC_MAPPING_NAMES, MODEL_TYPES[:EncoderOnly]], [MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES, MODEL_TYPES[:EncoderOnly]], [MODEL_FOR_TEXT_TO_SPECTROGRAM_MAPPING_NAMES, MODEL_TYPES[:Seq2Seq]], [MODEL_FOR_TEXT_TO_WAVEFORM_MAPPING_NAMES, MODEL_TYPES[:EncoderOnly]], [MODEL_FOR_AUDIO_XVECTOR_MAPPING_NAMES, MODEL_TYPES[:EncoderOnly]], [MODEL_FOR_AUDIO_FRAME_CLASSIFICATION_MAPPING_NAMES, MODEL_TYPES[:EncoderOnly]], [MODEL_FOR_IMAGE_FEATURE_EXTRACTION_MAPPING_NAMES, MODEL_TYPES[:EncoderOnly]] ]
- VERSION =
"1.1.1"
- SUPPORTED_TASKS =
{ "text-classification" => { tokenizer: AutoTokenizer, pipeline: TextClassificationPipeline, model: AutoModelForSequenceClassification, default: { model: "Xenova/distilbert-base-uncased-finetuned-sst-2-english" }, type: "text" }, "token-classification" => { tokenizer: AutoTokenizer, pipeline: TokenClassificationPipeline, model: AutoModelForTokenClassification, default: { model: "Xenova/bert-base-multilingual-cased-ner-hrl" }, type: "text" }, "question-answering" => { tokenizer: AutoTokenizer, pipeline: QuestionAnsweringPipeline, model: AutoModelForQuestionAnswering, default: { model: "Xenova/distilbert-base-cased-distilled-squad" }, type: "text" }, "fill-mask" => { tokenizer: AutoTokenizer, pipeline: FillMaskPipeline, model: AutoModelForMaskedLM, default: { model: "Xenova/bert-base-uncased" }, type: "text" }, "summarization" => { tokenizer: AutoTokenizer, pipeline: SummarizationPipeline, model: AutoModelForSeq2SeqLM, default: { model: "Xenova/distilbart-cnn-6-6" }, type: "text" }, "translation" => { tokenizer: AutoTokenizer, pipeline: TranslationPipeline, model: AutoModelForSeq2SeqLM, default: { model: "Xenova/t5-small" }, type: "text" }, "text2text-generation" => { tokenizer: AutoTokenizer, pipeline: Text2TextGenerationPipeline, model: AutoModelForSeq2SeqLM, default: { model: "Xenova/flan-t5-small" }, type: "text" }, "text-generation" => { tokenizer: AutoTokenizer, pipeline: TextGenerationPipeline, model: AutoModelForCausalLM, default: { model: "Xenova/gpt2" }, type: "text" }, "zero-shot-classification" => { tokenizer: AutoTokenizer, pipeline: ZeroShotClassificationPipeline, model: AutoModelForSequenceClassification, default: { model: "Xenova/distilbert-base-uncased-mnli" }, type: "text" }, "audio-classification" => { pipeline: AudioClassificationPipeline, model: AutoModelForAudioClassification, processor: AutoProcessor, default: { model: "Xenova/wav2vec2-base-superb-ks" }, type: "audio" }, # TODO # "zero-shot-audio-classification" => { # tokenizer: AutoTokenizer, # pipeline: ZeroShotAudioClassificationPipeline, # model: AutoModel, # processor: AutoProcessor, # default: { # model: "Xenova/clap-htsat-unfused" # }, # type: "multimodal" # }, # TODO # "automatic-speech-recognition" => { # tokenizer: AutoTokenizer, # pipeline: AutomaticSpeechRecognitionPipeline, # model: [AutoModelForSpeechSeq2Seq, AutoModelForCTC], # processor: AutoProcessor, # default: { # model: "Xenova/whisper-tiny.en" # }, # type: "multimodal" # }, "text-to-audio" => { tokenizer: AutoTokenizer, pipeline: TextToAudioPipeline, model: [AutoModelForTextToWaveform, AutoModelForTextToSpectrogram], processor: [AutoProcessor, nil], default: { model: "Xenova/speecht5_tts" }, type: "text" }, "image-to-text" => { tokenizer: AutoTokenizer, pipeline: ImageToTextPipeline, model: AutoModelForVision2Seq, processor: AutoProcessor, default: { model: "Xenova/vit-gpt2-image-captioning" }, type: "multimodal" }, "image-classification" => { pipeline: ImageClassificationPipeline, model: AutoModelForImageClassification, processor: AutoProcessor, default: { model: "Xenova/vit-base-patch16-224" }, type: "multimodal" }, "image-segmentation" => { pipeline: ImageSegmentationPipeline, model: [AutoModelForImageSegmentation, AutoModelForSemanticSegmentation], processor: AutoProcessor, default: { model: "Xenova/detr-resnet-50-panoptic" }, type: "multimodal" }, "zero-shot-image-classification" => { tokenizer: AutoTokenizer, pipeline: ZeroShotImageClassificationPipeline, model: AutoModel, processor: AutoProcessor, default: { model: "Xenova/clip-vit-base-patch32" }, type: "multimodal" }, "object-detection" => { pipeline: ObjectDetectionPipeline, model: AutoModelForObjectDetection, processor: AutoProcessor, default: { model: "Xenova/detr-resnet-50" }, type: "multimodal" }, "zero-shot-object-detection" => { tokenizer: AutoTokenizer, pipeline: ZeroShotObjectDetectionPipeline, model: AutoModelForZeroShotObjectDetection, processor: AutoProcessor, default: { model: "Xenova/owlvit-base-patch32" }, type: "multimodal" }, "document-question-answering" => { tokenizer: AutoTokenizer, pipeline: DocumentQuestionAnsweringPipeline, model: AutoModelForDocumentQuestionAnswering, processor: AutoProcessor, default: { model: "Xenova/donut-base-finetuned-docvqa" }, type: "multimodal" }, "image-to-image" => { pipeline: ImageToImagePipeline, model: AutoModelForImageToImage, processor: AutoProcessor, default: { model: "Xenova/swin2SR-classical-sr-x2-64" }, type: "image" }, "depth-estimation" => { pipeline: DepthEstimationPipeline, model: AutoModelForDepthEstimation, processor: AutoProcessor, default: { model: "Xenova/dpt-large" }, type: "image" }, "feature-extraction" => { tokenizer: AutoTokenizer, pipeline: FeatureExtractionPipeline, model: AutoModel, default: { model: "Xenova/all-MiniLM-L6-v2" }, type: "text" }, "image-feature-extraction" => { processor: AutoProcessor, pipeline: ImageFeatureExtractionPipeline, model: [AutoModelForImageFeatureExtraction, AutoModel], default: { model: "Xenova/vit-base-patch16-224" }, type: "image" }, "embedding" => { tokenizer: AutoTokenizer, pipeline: EmbeddingPipeline, model: AutoModel, default: { model: "sentence-transformers/all-MiniLM-L6-v2" }, type: "text" }, "reranking" => { tokenizer: AutoTokenizer, pipeline: RerankingPipeline, model: AutoModel, default: { model: "mixedbread-ai/mxbai-rerank-base-v1" }, type: "text" } }
- TASK_ALIASES =
{ "sentiment-analysis" => "text-classification", "ner" => "token-classification", "text-to-speech" => "text-to-audio" }
- DEFAULT_PROGRESS_CALLBACK =
lambda do |msg| stream = $stderr tty = stream.tty? width = tty ? stream.winsize[1] : 80 if msg[:status] == "progress" && tty stream.print "\r#{Utils::Hub.display_progress(msg[:file], width, msg[:size], msg[:total_size])}" elsif msg[:status] == "done" && !msg[:cache_hit] if tty stream.puts else stream.puts Utils::Hub.display_progress(msg[:file], width, 1, 1) end end end
- NO_DEFAULT =
Object.new
Class Attribute Summary collapse
-
.allow_remote_models ⇒ Object
Returns the value of attribute allow_remote_models.
-
.cache_dir ⇒ Object
Returns the value of attribute cache_dir.
-
.remote_host ⇒ Object
Returns the value of attribute remote_host.
-
.remote_path_template ⇒ Object
Returns the value of attribute remote_path_template.
Class Method Summary collapse
Class Attribute Details
.allow_remote_models ⇒ Object
Returns the value of attribute allow_remote_models.
6 7 8 |
# File 'lib/informers/env.rb', line 6 def allow_remote_models @allow_remote_models end |
.cache_dir ⇒ Object
Returns the value of attribute cache_dir.
6 7 8 |
# File 'lib/informers/env.rb', line 6 def cache_dir @cache_dir end |
.remote_host ⇒ Object
Returns the value of attribute remote_host.
6 7 8 |
# File 'lib/informers/env.rb', line 6 def remote_host @remote_host end |
.remote_path_template ⇒ Object
Returns the value of attribute remote_path_template.
6 7 8 |
# File 'lib/informers/env.rb', line 6 def remote_path_template @remote_path_template end |
Class Method Details
.pipeline(task, model = nil, quantized: NO_DEFAULT, progress_callback: DEFAULT_PROGRESS_CALLBACK, config: nil, cache_dir: nil, local_files_only: false, revision: "main", model_file_name: nil) ⇒ Object
1354 1355 1356 1357 1358 1359 1360 1361 1362 1363 1364 1365 1366 1367 1368 1369 1370 1371 1372 1373 1374 1375 1376 1377 1378 1379 1380 1381 1382 1383 1384 1385 1386 1387 1388 1389 1390 1391 1392 1393 1394 1395 1396 1397 1398 1399 1400 1401 1402 1403 1404 1405 1406 1407 1408 1409 1410 1411 1412 1413 1414 1415 1416 1417 1418 |
# File 'lib/informers/pipelines.rb', line 1354 def pipeline( task, model = nil, quantized: NO_DEFAULT, progress_callback: DEFAULT_PROGRESS_CALLBACK, config: nil, cache_dir: nil, local_files_only: false, revision: "main", model_file_name: nil ) # Apply aliases task = TASK_ALIASES[task] || task if quantized == NO_DEFAULT # TODO move default to task class quantized = ["text-classification", "token-classification", "question-answering", "feature-extraction"].include?(task) end # Get pipeline info pipeline_info = SUPPORTED_TASKS[task.split("_", 1)[0]] if !pipeline_info raise Error, "Unsupported pipeline: #{task}. Must be one of #{SUPPORTED_TASKS.keys}" end # Use model if specified, otherwise, use default if !model model = pipeline_info[:default][:model] warn "No model specified. Using default model: #{model.inspect}." end = { quantized:, progress_callback:, config:, cache_dir:, local_files_only:, revision:, model_file_name: } classes = { tokenizer: pipeline_info[:tokenizer], model: pipeline_info[:model], processor: pipeline_info[:processor] } # Load model, tokenizer, and processor (if they exist) results = load_items(classes, model, ) results[:task] = task # for previous revision of sentence-transformers/all-MiniLM-L6-v2 if model == "sentence-transformers/all-MiniLM-L6-v2" && results[:model].instance_variable_get(:@session).outputs.any? { |v| v[:name] == "token_embeddings" } results[:model].instance_variable_set(:@output_names, ["token_embeddings"]) end Utils.dispatch_callback(progress_callback, { status: "ready", task: task, model: model }) pipeline_class = pipeline_info.fetch(:pipeline) pipeline_class.new(**results) end |