def infer_task_from_hub()

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

• 10 lines of code
• 2 McCabe index (conditional complexity)

def infer_task_from_hub(model_id: str, revision: Optional[str] = None, use_auth_token: Optional[str] = None) -> str:
    """
    Infer task from Hub by extracting `pipeline_tag` for model_info.
    """
    _api = HfApi()
    model_info = _api.model_info(repo_id=model_id, revision=revision, token=use_auth_token)
    if model_info.pipeline_tag is not None:
        # set env to work with
        os.environ["HF_TASK"] = model_info.pipeline_tag
        return model_info.pipeline_tag
    else:
        raise ValueError(
            f"Task couldn't be inferenced from {model_info.pipeline_tag}." "Use env `HF_TASK` to define your task."
        )