def get_pipeline()

in src/sagemaker_huggingface_inference_toolkit/transformers_utils.py [0:0]

• 34 lines of code
• 12 McCabe index (conditional complexity)

def get_pipeline(task: str, device: int, model_dir: Path, **kwargs) -> Pipeline:
    """
    create pipeline class for a specific task based on local saved model
    """
    if task is None:
        raise EnvironmentError(
            "The task for this model is not set: Please set one: https://huggingface.co/docs#how-is-a-models-type-of-inference-api-and-widget-determined"
        )
    # define tokenizer or feature extractor as kwargs to load it the pipeline correctly
    if task in {
        "automatic-speech-recognition",
        "image-segmentation",
        "image-classification",
        "audio-classification",
        "object-detection",
        "zero-shot-image-classification",
    }:
        kwargs["feature_extractor"] = model_dir
    else:
        kwargs["tokenizer"] = model_dir
    # check if optimum neuron is available and tries to load it
    if is_optimum_neuron_available():
        hf_pipeline = get_optimum_neuron_pipeline(task=task, model_dir=model_dir)
    elif TRUST_REMOTE_CODE and os.environ.get("HF_MODEL_ID", None) is not None and device == 0:
        tokenizer = AutoTokenizer.from_pretrained(os.environ["HF_MODEL_ID"])

        hf_pipeline = pipeline(
            task=task,
            model=os.environ["HF_MODEL_ID"],
            tokenizer=tokenizer,
            trust_remote_code=TRUST_REMOTE_CODE,
            model_kwargs={"device_map": "auto", "torch_dtype": "auto"},
        )
    elif is_diffusers_available() and task == "text-to-image":
        hf_pipeline = get_diffusers_pipeline(task=task, model_dir=model_dir, device=device, **kwargs)
    else:
        # load pipeline
        hf_pipeline = pipeline(
            task=task, model=model_dir, device=device, trust_remote_code=TRUST_REMOTE_CODE, **kwargs
        )

    return hf_pipeline