mlflow--mlflow/tests/transformers/test_transformers_model_export.py

import base64
import gc
import importlib.util
import json
import math
import os
import pathlib
import re
import shutil
import textwrap
from pathlib import Path
from unittest import mock

import huggingface_hub
import librosa
import numpy as np
import pandas as pd
import pytest
import torch
import transformers
import yaml
from datasets import load_dataset
from huggingface_hub import ModelCard
from packaging.version import Version

import mlflow
import mlflow.pyfunc.scoring_server as pyfunc_scoring_server
from mlflow import pyfunc
from mlflow.deployments import PredictionsResponse
from mlflow.exceptions import MlflowException
from mlflow.models import Model, ModelSignature, infer_signature
from mlflow.models.model import METADATA_FILES
from mlflow.store.artifact.s3_artifact_repo import S3ArtifactRepository
from mlflow.tracking.artifact_utils import _download_artifact_from_uri
from mlflow.transformers import (
    _CARD_DATA_FILE_NAME,
    _CARD_TEXT_FILE_NAME,
    _build_pipeline_from_model_input,
    _fetch_model_card,
    _get_task_for_model,
    _is_model_distributed_in_memory,
    _should_add_pyfunc_to_model,
    _TransformersWrapper,
    _try_import_conversational_pipeline,
    _validate_llm_inference_task_type,
    _write_card_data,
    _write_license_information,
    get_default_conda_env,
    get_default_pip_requirements,
)
from mlflow.types.schema import Array, ColSpec, DataType, ParamSchema, ParamSpec, Schema
from mlflow.utils.environment import _mlflow_conda_env

from tests.helper_functions import (
    _assert_pip_requirements,
    _compare_conda_env_requirements,
    _compare_logged_code_paths,
    _get_deps_from_requirement_file,
    _mlflow_major_version_string,
    assert_register_model_called_with_local_model_path,
    flaky,
    pyfunc_scoring_endpoint,
    pyfunc_serve_and_score_model,
)
from tests.transformers.helper import (
    CHAT_TEMPLATE,
    IS_NEW_FEATURE_EXTRACTION_API,
    IS_TRANSFORMERS_V5_OR_LATER,
)
from tests.transformers.test_transformers_peft_model import SKIP_IF_PEFT_NOT_AVAILABLE

# NB: Some pipelines under test in this suite come very close or outright exceed the
# default runner containers specs of 7GB RAM. Due to this inability to run the suite without
# generating a SIGTERM Error (143), some tests are marked as local only.
# See: https://docs.github.com/en/actions/using-github-hosted-runners/about-github-hosted- \
# runners#supported-runners-and-hardware-resources for instance specs.
RUNNING_IN_GITHUB_ACTIONS = os.environ.get("GITHUB_ACTIONS") == "true"
GITHUB_ACTIONS_SKIP_REASON = "Test consumes too much memory"

skip_transformers_v5_or_later = pytest.mark.skipif(
    IS_TRANSFORMERS_V5_OR_LATER,
    reason="Incompatible API changes in transformers 5.x",
)
image_url = "https://raw.githubusercontent.com/mlflow/mlflow/master/tests/datasets/cat.png"
image_file_path = pathlib.Path(pathlib.Path(__file__).parent.parent, "datasets", "cat.png")
# Test that can only be run locally:
# - Summarization pipeline tests
# - TextClassifier pipeline tests
# - Text2TextGeneration pipeline tests
# - Conversational pipeline tests


@pytest.fixture(autouse=True)
def force_gc():
    # This reduces the memory pressure for the usage of the larger pipeline fixtures ~500MB - 1GB
    gc.disable()
    gc.collect()
    gc.set_threshold(0)
    gc.collect()
    gc.enable()


@pytest.fixture
def model_path(tmp_path):
    model_path = tmp_path.joinpath("model")
    yield model_path

    # Pytest keeps the temporary directory created by `tmp_path` fixture for 3 recent test sessions
    # by default. This is useful for debugging during local testing, but in CI it just wastes the
    # disk space.
    if os.environ.get("GITHUB_ACTIONS") == "true":
        shutil.rmtree(model_path, ignore_errors=True)


@pytest.fixture
def transformers_custom_env(tmp_path):
    conda_env = tmp_path.joinpath("conda_env.yml")
    _mlflow_conda_env(conda_env, additional_pip_deps=["transformers"])
    return conda_env


@pytest.fixture
def mock_pyfunc_wrapper():
    return mlflow.transformers._TransformersWrapper("mock")


@pytest.fixture
@flaky()
def image_for_test():
    dataset = load_dataset("hf-internal-testing/dummy_image_text_data")
    return dataset["train"]["image"][3]


@pytest.mark.parametrize(
    ("pipeline", "expected_requirements"),
    [
        ("small_qa_pipeline", {"transformers", "torch", "torchvision"}),
        pytest.param(
            "peft_pipeline",
            {"peft", "transformers", "torch", "torchvision"},
            marks=SKIP_IF_PEFT_NOT_AVAILABLE,
        ),
    ],
)
def test_default_requirements(pipeline, expected_requirements, request):
    if "torch" in expected_requirements and importlib.util.find_spec("accelerate"):
        expected_requirements.add("accelerate")

    model = request.getfixturevalue(pipeline).model
    pip_requirements = get_default_pip_requirements(model)
    conda_requirements = get_default_conda_env(model)["dependencies"][2]["pip"]

    def _strip_requirements(requirements):
        return {req.split("==")[0] for req in requirements}

    assert _strip_requirements(pip_requirements) == expected_requirements
    assert _strip_requirements(conda_requirements) == (expected_requirements | {"mlflow"})


def test_inference_task_validation(small_qa_pipeline):
    with pytest.raises(
        MlflowException, match="The task provided is invalid. 'llm/v1/invalid' is not"
    ):
        _validate_llm_inference_task_type("llm/v1/invalid", "text-generation")
    with pytest.raises(
        MlflowException, match="The task provided is invalid. 'llm/v1/completions' is not"
    ):
        _validate_llm_inference_task_type("llm/v1/completions", small_qa_pipeline)
    _validate_llm_inference_task_type("llm/v1/completions", "text-generation")


@pytest.mark.parametrize(
    ("model", "result"),
    [
        ("small_qa_pipeline", True),
        ("small_multi_modal_pipeline", False),
        ("small_vision_model", True),
    ],
)
def test_pipeline_eligibility_for_pyfunc_registration(model, result, request):
    pipeline = request.getfixturevalue(model)
    assert _should_add_pyfunc_to_model(pipeline) == result


def test_component_multi_modal_model_ineligible_for_pyfunc(component_multi_modal):
    task = transformers.pipelines.get_task(component_multi_modal["model"].name_or_path)
    pipeline = _build_pipeline_from_model_input(component_multi_modal, task)
    assert not _should_add_pyfunc_to_model(pipeline)


def test_pipeline_construction_from_base_nlp_model(small_qa_pipeline):
    generated = _build_pipeline_from_model_input(
        {"model": small_qa_pipeline.model, "tokenizer": small_qa_pipeline.tokenizer},
        "question-answering",
    )
    assert isinstance(generated, type(small_qa_pipeline))
    assert isinstance(generated.tokenizer, type(small_qa_pipeline.tokenizer))


def test_pipeline_construction_from_base_vision_model(small_vision_model):
    model = {"model": small_vision_model.model, "tokenizer": small_vision_model.tokenizer}
    if IS_NEW_FEATURE_EXTRACTION_API:
        model.update({"image_processor": small_vision_model.image_processor})
    else:
        model.update({"feature_extractor": small_vision_model.feature_extractor})
    generated = _build_pipeline_from_model_input(model, task="image-classification")
    assert isinstance(generated, type(small_vision_model))
    assert isinstance(generated.tokenizer, type(small_vision_model.tokenizer))
    if IS_NEW_FEATURE_EXTRACTION_API:
        assert isinstance(generated.image_processor, type(small_vision_model.image_processor))
    else:
        assert isinstance(generated.feature_extractor, transformers.MobileNetV2ImageProcessor)


def test_saving_with_invalid_dict_as_model(model_path):
    with pytest.raises(
        MlflowException, match="Invalid dictionary submitted for 'transformers_model'. The "
    ):
        mlflow.transformers.save_model(transformers_model={"invalid": "key"}, path=model_path)

    with pytest.raises(
        MlflowException, match="The 'transformers_model' dictionary must have an entry"
    ):
        mlflow.transformers.save_model(
            transformers_model={"tokenizer": "some_tokenizer"}, path=model_path
        )


def test_model_card_acquisition_vision_model(small_vision_model):
    model_provided_card = _fetch_model_card(small_vision_model.model.name_or_path)
    assert model_provided_card.data.to_dict()["tags"] == ["vision", "image-classification"]
    assert len(model_provided_card.text) > 0


@pytest.mark.parametrize(
    ("repo_id", "license_file"),
    [
        ("google/mobilenet_v2_1.0_224", "LICENSE.txt"),  # no license declared
        ("csarron/mobilebert-uncased-squad-v2", "LICENSE.txt"),  # mit license
        ("codellama/CodeLlama-34b-hf", "LICENSE"),  # custom license
        ("openai/whisper-tiny", "LICENSE.txt"),  # apache license
        ("stabilityai/stable-code-3b", "LICENSE"),  # custom
        ("mistralai/Mixtral-8x7B-Instruct-v0.1", "LICENSE.txt"),  # apache
    ],
)
def test_license_acquisition(repo_id, license_file, tmp_path):
    card_data = _fetch_model_card(repo_id)
    _write_license_information(repo_id, card_data, tmp_path)
    license_file = list(tmp_path.glob("*LICENSE*"))
    assert len(license_file) == 1
    assert tmp_path.joinpath(license_file[0]).stat().st_size > 0


def test_license_fallback(tmp_path):
    _write_license_information("not a real repo", None, tmp_path)
    assert tmp_path.joinpath("LICENSE.txt").stat().st_size > 0


def test_vision_model_save_pipeline_with_defaults(small_vision_model, model_path):
    mlflow.transformers.save_model(transformers_model=small_vision_model, path=model_path)
    # validate inferred pip requirements
    requirements = model_path.joinpath("requirements.txt").read_text()
    reqs = {req.split("==")[0] for req in requirements.split("\n")}
    expected_requirements = {"torch", "torchvision", "transformers"}
    assert reqs.intersection(expected_requirements) == expected_requirements
    # validate inferred model card data
    card_data = yaml.safe_load(model_path.joinpath("model_card_data.yaml").read_bytes())
    assert card_data["tags"] == ["vision", "image-classification"]
    # verify the license file has been written
    license_file = model_path.joinpath("LICENSE.txt").read_text()
    assert len(license_file) > 0
    # Validate inferred model card text
    with model_path.joinpath("model_card.md").open() as file:
        card_text = file.read()
    assert len(card_text) > 0
    # Validate conda.yaml
    conda_env = yaml.safe_load(model_path.joinpath("conda.yaml").read_bytes())
    assert {req.split("==")[0] for req in conda_env["dependencies"][2]["pip"]}.intersection(
        expected_requirements
    ) == expected_requirements
    # Validate the MLModel file
    mlmodel = yaml.safe_load(model_path.joinpath("MLmodel").read_bytes())
    flavor_config = mlmodel["flavors"]["transformers"]
    assert flavor_config["instance_type"] == "ImageClassificationPipeline"
    assert flavor_config["pipeline_model_type"] == "MobileNetV2ForImageClassification"
    assert flavor_config["task"] == "image-classification"
    assert flavor_config["source_model_name"] == "google/mobilenet_v2_1.0_224"


def test_vision_model_save_model_for_task_and_card_inference(small_vision_model, model_path):
    mlflow.transformers.save_model(transformers_model=small_vision_model, path=model_path)
    # validate inferred pip requirements
    requirements = model_path.joinpath("requirements.txt").read_text()
    reqs = {req.split("==")[0] for req in requirements.split("\n")}
    expected_requirements = {"torch", "torchvision", "transformers"}
    assert reqs.intersection(expected_requirements) == expected_requirements
    # validate inferred model card data
    card_data = yaml.safe_load(model_path.joinpath("model_card_data.yaml").read_bytes())
    assert card_data["tags"] == ["vision", "image-classification"]
    # Validate inferred model card text
    card_text = model_path.joinpath("model_card.md").read_text(encoding="utf-8")
    assert len(card_text) > 0
    # verify the license file has been written
    license_file = model_path.joinpath("LICENSE.txt").read_text()
    assert len(license_file) > 0
    # Validate the MLModel file
    mlmodel = yaml.safe_load(model_path.joinpath("MLmodel").read_bytes())
    flavor_config = mlmodel["flavors"]["transformers"]
    assert flavor_config["instance_type"] == "ImageClassificationPipeline"
    assert flavor_config["pipeline_model_type"] == "MobileNetV2ForImageClassification"
    assert flavor_config["task"] == "image-classification"
    assert flavor_config["source_model_name"] == "google/mobilenet_v2_1.0_224"


def test_qa_model_save_model_for_task_and_card_inference(small_qa_pipeline, model_path):
    mlflow.transformers.save_model(
        transformers_model={
            "model": small_qa_pipeline.model,
            "tokenizer": small_qa_pipeline.tokenizer,
        },
        path=model_path,
    )
    # validate inferred pip requirements
    with model_path.joinpath("requirements.txt").open() as file:
        requirements = file.read()
    reqs = {req.split("==")[0] for req in requirements.split("\n")}
    expected_requirements = {"torch", "transformers"}
    assert reqs.intersection(expected_requirements) == expected_requirements
    # validate that the card was acquired by model reference
    card_data = yaml.safe_load(model_path.joinpath("model_card_data.yaml").read_bytes())
    assert card_data["datasets"] == ["squad_v2"]
    assert "tags" in card_data
    # verify the license file has been written
    license_file = model_path.joinpath("LICENSE.txt").read_text()
    assert len(license_file) > 0
    # Validate inferred model card text
    with model_path.joinpath("model_card.md").open() as file:
        card_text = file.read()
    assert len(card_text) > 0
    # validate MLmodel files
    mlmodel = yaml.safe_load(model_path.joinpath("MLmodel").read_bytes())
    flavor_config = mlmodel["flavors"]["transformers"]
    assert flavor_config["instance_type"] == "QuestionAnsweringPipeline"
    assert flavor_config["pipeline_model_type"] == "MobileBertForQuestionAnswering"
    assert flavor_config["task"] == "question-answering"
    assert flavor_config["source_model_name"] == "csarron/mobilebert-uncased-squad-v2"


def test_qa_model_save_and_override_card(small_qa_pipeline, model_path):
    supplied_card = """
                    ---
                    language: en
                    license: bsd
                    ---

                    # I made a new model!
                    """
    card_info = textwrap.dedent(supplied_card)
    card = ModelCard(card_info)
    # save the model instance
    mlflow.transformers.save_model(
        transformers_model=small_qa_pipeline,
        path=model_path,
        model_card=card,
    )
    # validate that the card was acquired by model reference
    card_data = yaml.safe_load(model_path.joinpath("model_card_data.yaml").read_bytes())
    assert card_data["language"] == "en"
    assert card_data["license"] == "bsd"
    # Validate inferred model card text
    with model_path.joinpath("model_card.md").open() as file:
        card_text = file.read()
    # verify the license file has been written
    license_file = model_path.joinpath("LICENSE.txt").read_text()
    assert len(license_file) > 0
    assert card_text.startswith("\n# I made a new model!")
    # validate MLmodel files
    mlmodel = yaml.safe_load(model_path.joinpath("MLmodel").read_bytes())
    flavor_config = mlmodel["flavors"]["transformers"]
    assert flavor_config["instance_type"] == "QuestionAnsweringPipeline"
    assert flavor_config["pipeline_model_type"] == "MobileBertForQuestionAnswering"
    assert flavor_config["task"] == "question-answering"
    assert flavor_config["source_model_name"] == "csarron/mobilebert-uncased-squad-v2"


def test_basic_save_model_and_load_text_pipeline(text_classification_pipeline, model_path):
    mlflow.transformers.save_model(
        transformers_model={
            "model": text_classification_pipeline.model,
            "tokenizer": text_classification_pipeline.tokenizer,
        },
        path=model_path,
    )
    loaded = mlflow.transformers.load_model(model_path)
    result = loaded("MLflow is a really neat tool!")
    assert result[0]["label"] == "POSITIVE"
    assert result[0]["score"] > 0.5


@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float64])
def test_basic_save_model_with_torch_dtype(text2text_generation_pipeline, model_path, dtype):
    mlflow.transformers.save_model(
        transformers_model=text2text_generation_pipeline,
        path=model_path,
        torch_dtype=dtype,
    )

    loaded = mlflow.transformers.load_model(model_path)
    assert loaded.model.dtype == dtype

    loaded = mlflow.transformers.load_model(model_path, torch_dtype=torch.float32)
    assert loaded.model.dtype == torch.float32


def test_basic_save_model_and_load_vision_pipeline(small_vision_model, model_path, image_for_test):
    if IS_NEW_FEATURE_EXTRACTION_API:
        model = {
            "model": small_vision_model.model,
            "image_processor": small_vision_model.image_processor,
            "tokenizer": small_vision_model.tokenizer,
        }
    else:
        model = {
            "model": small_vision_model.model,
            "feature_extractor": small_vision_model.feature_extractor,
            "tokenizer": small_vision_model.tokenizer,
        }
    mlflow.transformers.save_model(
        transformers_model=model,
        path=model_path,
    )
    loaded = mlflow.transformers.load_model(model_path)
    prediction = loaded(image_for_test)
    assert prediction[0]["label"] == "wall clock"
    assert prediction[0]["score"] > 0.5


@flaky()
def test_multi_modal_pipeline_save_and_load(small_multi_modal_pipeline, model_path, image_for_test):
    mlflow.transformers.save_model(transformers_model=small_multi_modal_pipeline, path=model_path)
    question = "How many wall clocks are in the picture?"
    # Load components
    components = mlflow.transformers.load_model(model_path, return_type="components")
    if IS_NEW_FEATURE_EXTRACTION_API:
        expected_components = {"model", "task", "tokenizer", "image_processor"}
    else:
        expected_components = {"model", "task", "tokenizer", "feature_extractor"}
    assert set(components.keys()).intersection(expected_components) == expected_components
    constructed_pipeline = transformers.pipeline(**components)
    answer = constructed_pipeline(image=image_for_test, question=question)
    assert answer[0]["answer"] == "1"
    # Load pipeline
    pipeline = mlflow.transformers.load_model(model_path)
    pipeline_answer = pipeline(image=image_for_test, question=question)
    assert pipeline_answer[0]["answer"] == "1"
    # Test invalid loading mode
    with pytest.raises(MlflowException, match="The specified return_type mode 'magic' is"):
        mlflow.transformers.load_model(model_path, return_type="magic")


def test_multi_modal_component_save_and_load(component_multi_modal, model_path, image_for_test):
    if IS_NEW_FEATURE_EXTRACTION_API:
        processor = component_multi_modal["image_processor"]
    else:
        processor = component_multi_modal["feature_extractor"]
    mlflow.transformers.save_model(
        transformers_model=component_multi_modal,
        path=model_path,
        processor=processor,
    )
    # Ensure that the appropriate Processor object was detected and loaded with the pipeline.
    loaded_components = mlflow.transformers.load_model(
        model_uri=model_path, return_type="components"
    )
    assert isinstance(loaded_components["model"], transformers.ViltForQuestionAnswering)
    assert isinstance(loaded_components["tokenizer"], transformers.BertTokenizerFast)
    # This is to simulate a post-processing processor that would be used externally to a Pipeline
    # This isn't being tested on an actual use case of such a model type due to the size of
    # these types of models that have this interface being ill-suited for CI testing.

    if IS_NEW_FEATURE_EXTRACTION_API:
        processor_key = "image_processor"
        assert isinstance(loaded_components[processor_key], transformers.ViltImageProcessor)
    else:
        processor_key = "feature_extractor"
        assert isinstance(loaded_components[processor_key], transformers.ViltProcessor)
        assert isinstance(loaded_components["processor"], transformers.ViltProcessor)
    if not IS_NEW_FEATURE_EXTRACTION_API:
        # NB: This simulated behavior is no longer valid in versions 4.27.4 and above.
        # With the port of functionality away from feature extractor types, the new architecture
        # for multi-modal models is entirely pipeline based.
        # Make sure that the component usage works correctly when extracted from inference loading
        model = loaded_components["model"]
        processor = loaded_components["processor"]
        question = "What are the cats doing?"
        inputs = processor(image_for_test, question, return_tensors="pt")
        outputs = model(**inputs)
        logits = outputs.logits
        idx = logits.argmax(-1).item()
        answer = model.config.id2label[idx]
        assert answer == "sleeping"


@flaky()
def test_pipeline_saved_model_with_processor_cannot_be_loaded_as_pipeline(
    component_multi_modal, model_path
):
    invalid_pipeline = transformers.pipeline(
        task="visual-question-answering", **component_multi_modal
    )
    if IS_NEW_FEATURE_EXTRACTION_API:
        processor = component_multi_modal["image_processor"]
    else:
        processor = component_multi_modal["feature_extractor"]
    mlflow.transformers.save_model(
        transformers_model=invalid_pipeline,
        path=model_path,
        processor=processor,  # If this is specified, we cannot guarantee correct inference
    )
    with pytest.raises(
        MlflowException, match="This model has been saved with a processor. Processor objects"
    ):
        mlflow.transformers.load_model(model_uri=model_path, return_type="pipeline")


def test_component_saved_model_with_processor_cannot_be_loaded_as_pipeline(
    component_multi_modal, model_path
):
    if IS_NEW_FEATURE_EXTRACTION_API:
        processor = component_multi_modal["image_processor"]
    else:
        processor = component_multi_modal["feature_extractor"]
    mlflow.transformers.save_model(
        transformers_model=component_multi_modal,
        path=model_path,
        processor=processor,
    )
    with pytest.raises(
        MlflowException,
        match="This model has been saved with a processor. Processor objects are not compatible "
        "with Pipelines. Please load",
    ):
        mlflow.transformers.load_model(model_uri=model_path, return_type="pipeline")


@pytest.mark.parametrize("should_start_run", [True, False])
def test_log_and_load_transformers_pipeline(small_qa_pipeline, tmp_path, should_start_run):
    try:
        if should_start_run:
            mlflow.start_run()
        artifact_path = "transformers"
        conda_env = tmp_path.joinpath("conda_env.yaml")
        _mlflow_conda_env(conda_env, additional_pip_deps=["transformers"])
        model_info = mlflow.transformers.log_model(
            small_qa_pipeline,
            name=artifact_path,
            conda_env=str(conda_env),
        )
        reloaded_model = mlflow.transformers.load_model(
            model_uri=model_info.model_uri, return_type="pipeline"
        )
        assert (
            reloaded_model(
                question="Who's house?", context="The house is owned by a man named Run."
            )["answer"]
            == "Run"
        )
        model_path = pathlib.Path(_download_artifact_from_uri(artifact_uri=model_info.model_uri))
        model_config = Model.load(str(model_path.joinpath("MLmodel")))
        assert pyfunc.FLAVOR_NAME in model_config.flavors
        assert pyfunc.ENV in model_config.flavors[pyfunc.FLAVOR_NAME]
        env_path = model_config.flavors[pyfunc.FLAVOR_NAME][pyfunc.ENV]["conda"]
        assert model_path.joinpath(env_path).exists()
    finally:
        mlflow.end_run()


def test_load_pipeline_from_remote_uri_succeeds(
    text_classification_pipeline, model_path, mock_s3_bucket
):
    mlflow.transformers.save_model(transformers_model=text_classification_pipeline, path=model_path)
    artifact_root = f"s3://{mock_s3_bucket}"
    artifact_path = "model"
    artifact_repo = S3ArtifactRepository(artifact_root)
    artifact_repo.log_artifacts(model_path, artifact_path=artifact_path)
    model_uri = os.path.join(artifact_root, artifact_path)
    loaded = mlflow.transformers.load_model(model_uri=str(model_uri), return_type="pipeline")
    assert loaded("I like it when CI checks pass and are never flaky!")[0]["label"] == "POSITIVE"


def test_transformers_log_model_calls_register_model(small_qa_pipeline, tmp_path):
    artifact_path = "transformers"
    register_model_patch = mock.patch("mlflow.tracking._model_registry.fluent._register_model")
    with mlflow.start_run(), register_model_patch:
        conda_env = tmp_path.joinpath("conda_env.yaml")
        _mlflow_conda_env(conda_env, additional_pip_deps=["transformers", "torch", "torchvision"])
        model_info = mlflow.transformers.log_model(
            small_qa_pipeline,
            name=artifact_path,
            conda_env=str(conda_env),
            registered_model_name="Question-Answering Model 1",
        )
        assert_register_model_called_with_local_model_path(
            register_model_mock=mlflow.tracking._model_registry.fluent._register_model,
            model_uri=model_info.model_uri,
            registered_model_name="Question-Answering Model 1",
        )


def test_transformers_log_model_with_no_registered_model_name(small_vision_model, tmp_path):
    if IS_NEW_FEATURE_EXTRACTION_API:
        model = {
            "model": small_vision_model.model,
            "image_processor": small_vision_model.image_processor,
            "tokenizer": small_vision_model.tokenizer,
        }
    else:
        model = {
            "model": small_vision_model.model,
            "feature_extractor": small_vision_model.feature_extractor,
            "tokenizer": small_vision_model.tokenizer,
        }

    artifact_path = "transformers"
    registered_model_patch = mock.patch("mlflow.tracking._model_registry.fluent._register_model")
    with mlflow.start_run(), registered_model_patch:
        conda_env = tmp_path.joinpath("conda_env.yaml")
        _mlflow_conda_env(conda_env, additional_pip_deps=["tensorflow", "transformers"])
        mlflow.transformers.log_model(
            model,
            name=artifact_path,
            conda_env=str(conda_env),
        )
        mlflow.tracking._model_registry.fluent._register_model.assert_not_called()


def test_transformers_log_model_with_prompt_template_sets_return_full_text_false(
    text_generation_pipeline,
):
    artifact_path = "text_generation_with_prompt_template"
    prompt_template = "User: {prompt}"

    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            text_generation_pipeline,
            name=artifact_path,
            prompt_template=prompt_template,
        )

    model_path = pathlib.Path(_download_artifact_from_uri(model_info.model_uri))
    mlmodel = Model.load(str(model_path.joinpath("MLmodel")))

    pyfunc_flavor = mlmodel.flavors["python_function"]
    config = pyfunc_flavor.get("config")

    assert config.get("return_full_text") is False


def test_transformers_save_persists_requirements_in_mlflow_directory(
    small_qa_pipeline, model_path, transformers_custom_env
):
    mlflow.transformers.save_model(
        transformers_model=small_qa_pipeline,
        path=model_path,
        conda_env=str(transformers_custom_env),
    )
    saved_pip_req_path = model_path.joinpath("requirements.txt")
    _compare_conda_env_requirements(transformers_custom_env, saved_pip_req_path)


def test_transformers_log_with_pip_requirements(small_multi_modal_pipeline, tmp_path):
    expected_mlflow_version = _mlflow_major_version_string()

    requirements_file = tmp_path.joinpath("requirements.txt")
    requirements_file.write_text("coolpackage")
    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            small_multi_modal_pipeline, name="model", pip_requirements=str(requirements_file)
        )
        _assert_pip_requirements(
            model_info.model_uri, [expected_mlflow_version, "coolpackage"], strict=True
        )
    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            small_multi_modal_pipeline,
            name="model",
            pip_requirements=[f"-r {requirements_file}", "alsocool"],
        )
        _assert_pip_requirements(
            model_info.model_uri,
            [expected_mlflow_version, "coolpackage", "alsocool"],
            strict=True,
        )
    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            small_multi_modal_pipeline,
            name="model",
            pip_requirements=[f"-c {requirements_file}", "constrainedcool"],
        )
        _assert_pip_requirements(
            model_info.model_uri,
            [expected_mlflow_version, "constrainedcool", "-c constraints.txt"],
            ["coolpackage"],
            strict=True,
        )


def test_transformers_log_with_extra_pip_requirements(small_multi_modal_pipeline, tmp_path):
    expected_mlflow_version = _mlflow_major_version_string()
    default_requirements = mlflow.transformers.get_default_pip_requirements(
        small_multi_modal_pipeline.model
    )
    requirements_file = tmp_path.joinpath("requirements.txt")
    requirements_file.write_text("coolpackage")
    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            small_multi_modal_pipeline, name="model", extra_pip_requirements=str(requirements_file)
        )
        _assert_pip_requirements(
            model_info.model_uri,
            [expected_mlflow_version, *default_requirements, "coolpackage"],
            strict=True,
        )
    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            small_multi_modal_pipeline,
            name="model",
            extra_pip_requirements=[f"-r {requirements_file}", "alsocool"],
        )
        _assert_pip_requirements(
            model_info.model_uri,
            [expected_mlflow_version, *default_requirements, "coolpackage", "alsocool"],
            strict=True,
        )
    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            small_multi_modal_pipeline,
            name="model",
            extra_pip_requirements=[f"-c {requirements_file}", "constrainedcool"],
        )
        _assert_pip_requirements(
            model_info.model_uri,
            [
                expected_mlflow_version,
                *default_requirements,
                "constrainedcool",
                "-c constraints.txt",
            ],
            ["coolpackage"],
            strict=True,
        )


def test_transformers_log_with_duplicate_extra_pip_requirements(small_multi_modal_pipeline):
    with pytest.raises(
        MlflowException, match="The specified requirements versions are incompatible"
    ):
        with mlflow.start_run():
            mlflow.transformers.log_model(
                small_multi_modal_pipeline,
                name="model",
                extra_pip_requirements=["transformers==1.1.0"],
            )


def test_transformers_pt_model_save_without_conda_env_uses_default_env_with_expected_dependencies(
    small_qa_pipeline, model_path
):
    mlflow.transformers.save_model(small_qa_pipeline, model_path)
    _assert_pip_requirements(
        model_path, mlflow.transformers.get_default_pip_requirements(small_qa_pipeline.model)
    )
    pip_requirements = _get_deps_from_requirement_file(model_path)
    assert "tensorflow" not in pip_requirements
    assert "accelerate" in pip_requirements
    assert "torch" in pip_requirements


@pytest.mark.skipif(
    importlib.util.find_spec("accelerate") is not None, reason="fails when accelerate is installed"
)
def test_transformers_pt_model_save_dependencies_without_accelerate(
    text_generation_pipeline, model_path
):
    mlflow.transformers.save_model(text_generation_pipeline, model_path)
    _assert_pip_requirements(
        model_path, mlflow.transformers.get_default_pip_requirements(text_generation_pipeline.model)
    )
    pip_requirements = _get_deps_from_requirement_file(model_path)
    assert "tensorflow" not in pip_requirements
    assert "accelerate" not in pip_requirements
    assert "torch" in pip_requirements


def test_transformers_pt_model_log_without_conda_env_uses_default_env_with_expected_dependencies(
    small_qa_pipeline,
):
    artifact_path = "model"
    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(small_qa_pipeline, name=artifact_path)
    _assert_pip_requirements(
        model_info.model_uri,
        mlflow.transformers.get_default_pip_requirements(small_qa_pipeline.model),
    )
    pip_requirements = _get_deps_from_requirement_file(model_info.model_uri)
    assert "tensorflow" not in pip_requirements
    assert "torch" in pip_requirements


def test_log_model_with_code_paths(small_qa_pipeline):
    artifact_path = "model"
    with (
        mlflow.start_run(),
        mock.patch("mlflow.transformers._add_code_from_conf_to_system_path") as add_mock,
    ):
        model_info = mlflow.transformers.log_model(
            small_qa_pipeline, name=artifact_path, code_paths=[__file__]
        )
        model_uri = model_info.model_uri
        _compare_logged_code_paths(__file__, model_uri, mlflow.transformers.FLAVOR_NAME)
        mlflow.transformers.load_model(model_uri)
        add_mock.assert_called()


def test_non_existent_model_card_entry(small_qa_pipeline, model_path):
    with mock.patch("mlflow.transformers._fetch_model_card", return_value=None):
        mlflow.transformers.save_model(transformers_model=small_qa_pipeline, path=model_path)

        contents = {item.name for item in model_path.iterdir()}
        assert not contents.intersection({"model_card.txt", "model_card_data.yaml"})


def test_huggingface_hub_not_installed(small_qa_pipeline, model_path):
    with mock.patch.dict("sys.modules", {"huggingface_hub": None}):
        result = mlflow.transformers._fetch_model_card(small_qa_pipeline.model.name_or_path)

        assert result is None

        mlflow.transformers.save_model(transformers_model=small_qa_pipeline, path=model_path)

        contents = {item.name for item in model_path.iterdir()}
        assert not contents.intersection({"model_card.txt", "model_card_data.yaml"})

        license_data = model_path.joinpath("LICENSE.txt").read_text()
        assert license_data.rstrip().endswith("mobilebert-uncased-squad-v2")


@pytest.mark.skipif(
    _try_import_conversational_pipeline() is None,
    reason="Conversation model is deprecated and removed.",
)
def test_save_pipeline_without_defined_components(small_conversational_model, model_path):
    # This pipeline type explicitly does not have a configuration for an image_processor
    with mlflow.start_run():
        mlflow.transformers.save_model(
            transformers_model=small_conversational_model, path=model_path
        )
    pipe = mlflow.transformers.load_model(model_path)
    convo = transformers.Conversation("How are you today?")
    convo = pipe(convo)
    assert convo.generated_responses[-1] == "good"


@flaky()
def test_invalid_model_type_without_registered_name_does_not_save(model_path):
    invalid_pipeline = transformers.pipeline(task="text-generation", model="gpt2")
    del invalid_pipeline.model.name_or_path

    with pytest.raises(MlflowException, match="The submitted model type"):
        mlflow.transformers.save_model(transformers_model=invalid_pipeline, path=model_path)


def test_invalid_input_to_pyfunc_signature_output_wrapper_raises(component_multi_modal):
    with pytest.raises(MlflowException, match="The pipeline type submitted is not a valid"):
        mlflow.transformers.generate_signature_output(component_multi_modal["model"], "bogus")


@pytest.mark.parametrize(
    "inference_payload",
    [
        ({"question": "Who's house?", "context": "The house is owned by a man named Run."}),
        ([
            {
                "question": "What color is it?",
                "context": "Some people said it was green but I know that it's definitely blue",
            },
            {
                "question": "How do the wheels go?",
                "context": "The wheels on the bus go round and round. Round and round.",
            },
        ]),
        ([
            {
                "question": "What color is it?",
                "context": "Some people said it was green but I know that it's pink.",
            },
            {
                "context": "The people on the bus go up and down. Up and down.",
                "question": "How do the people go?",
            },
        ]),
    ],
)
def test_qa_pipeline_pyfunc_load_and_infer(small_qa_pipeline, model_path, inference_payload):
    signature = infer_signature(
        inference_payload,
        mlflow.transformers.generate_signature_output(small_qa_pipeline, inference_payload),
    )

    mlflow.transformers.save_model(
        transformers_model=small_qa_pipeline,
        path=model_path,
        signature=signature,
    )
    pyfunc_loaded = mlflow.pyfunc.load_model(model_path)

    inference = pyfunc_loaded.predict(inference_payload)

    assert isinstance(inference, list)
    assert all(isinstance(element, str) for element in inference)

    pd_input = (
        pd.DataFrame([inference_payload])
        if isinstance(inference_payload, dict)
        else pd.DataFrame(inference_payload)
    )
    pd_inference = pyfunc_loaded.predict(pd_input)

    assert isinstance(pd_inference, list)
    assert all(isinstance(element, str) for element in inference)


@pytest.mark.parametrize(
    "inference_payload",
    [
        image_url,
        str(image_file_path),
        "base64",
        pytest.param(
            "base64_encodebytes",
            marks=pytest.mark.skipif(
                Version(transformers.__version__) < Version("4.41"),
                reason="base64 encodebytes feature not present",
            ),
        ),
    ],
)
def test_vision_pipeline_pyfunc_load_and_infer(small_vision_model, model_path, inference_payload):
    if inference_payload == "base64":
        inference_payload = base64.b64encode(image_file_path.read_bytes()).decode("utf-8")
    elif inference_payload == "base64_encodebytes":
        inference_payload = base64.encodebytes(image_file_path.read_bytes()).decode("utf-8")
    signature = infer_signature(
        inference_payload,
        mlflow.transformers.generate_signature_output(small_vision_model, inference_payload),
    )
    mlflow.transformers.save_model(
        transformers_model=small_vision_model,
        path=model_path,
        signature=signature,
    )
    pyfunc_loaded = mlflow.pyfunc.load_model(model_path)
    predictions = pyfunc_loaded.predict(inference_payload)

    transformers_loaded_model = mlflow.transformers.load_model(model_path)
    expected_predictions = transformers_loaded_model.predict(inference_payload)
    assert list(predictions.to_dict("records")[0].values()) == expected_predictions


@pytest.mark.parametrize(
    ("data", "result"),
    [
        ("muppet keyboard type", ["A man is typing a muppet on a keyboard."]),
        (
            ["pencil draw paper", "pie apple eat"],
            # NB: The result of this test case, without inference config overrides is:
            # ["A man drawing on paper with pencil", "A man eating a pie with applies"]
            # The inference config override forces additional insertion of more grammatically
            # correct responses to validate that the inference config is being applied.
            ["A man draws a pencil on a paper.", "A man eats a pie of apples."],
        ),
    ],
)
def test_text2text_generation_pipeline_with_model_configs(
    text2text_generation_pipeline, tmp_path, data, result
):
    signature = infer_signature(
        data, mlflow.transformers.generate_signature_output(text2text_generation_pipeline, data)
    )

    model_config = {
        "top_k": 2,
        "num_beams": 5,
        "max_length": 30,
        "temperature": 0.62,
        "top_p": 0.85,
        "repetition_penalty": 1.15,
    }
    model_path1 = tmp_path.joinpath("model1")
    mlflow.transformers.save_model(
        text2text_generation_pipeline,
        path=model_path1,
        model_config=model_config,
        signature=signature,
    )
    pyfunc_loaded = mlflow.pyfunc.load_model(model_path1)

    inference = pyfunc_loaded.predict(data)

    assert inference == result

    pd_input = pd.DataFrame([data]) if isinstance(data, str) else pd.DataFrame(data)
    pd_inference = pyfunc_loaded.predict(pd_input)
    assert pd_inference == result

    model_path2 = tmp_path.joinpath("model2")
    signature_with_params = infer_signature(
        data,
        mlflow.transformers.generate_signature_output(text2text_generation_pipeline, data),
        model_config,
    )
    mlflow.transformers.save_model(
        text2text_generation_pipeline,
        path=model_path2,
        signature=signature_with_params,
    )
    pyfunc_loaded = mlflow.pyfunc.load_model(model_path2)

    dict_inference = pyfunc_loaded.predict(
        data,
        params=model_config,
    )

    assert dict_inference == inference


def test_text2text_generation_pipeline_with_model_config_and_params(
    text2text_generation_pipeline, model_path
):
    data = "muppet keyboard type"
    model_config = {
        "top_k": 2,
        "num_beams": 5,
        "top_p": 0.85,
        "repetition_penalty": 1.15,
        "do_sample": True,
    }
    parameters = {"top_k": 3, "max_new_tokens": 30}
    generated_output = mlflow.transformers.generate_signature_output(
        text2text_generation_pipeline, data
    )
    signature = infer_signature(
        data,
        generated_output,
        parameters,
    )

    mlflow.transformers.save_model(
        text2text_generation_pipeline,
        path=model_path,
        model_config=model_config,
        signature=signature,
    )
    pyfunc_loaded = mlflow.pyfunc.load_model(model_path)

    # model_config and default params are all applied
    res = pyfunc_loaded.predict(data)
    applied_params = model_config.copy()
    applied_params.update(parameters)
    res2 = pyfunc_loaded.predict(data, applied_params)
    assert res == res2

    assert res != pyfunc_loaded.predict(data, {"max_new_tokens": 3})

    # Extra params are ignored
    assert res == pyfunc_loaded.predict(data, {"extra_param": "extra_value"})


def test_text2text_generation_pipeline_with_params_success(
    text2text_generation_pipeline, model_path
):
    data = "muppet keyboard type"
    parameters = {"top_k": 2, "num_beams": 5, "do_sample": True}
    generated_output = mlflow.transformers.generate_signature_output(
        text2text_generation_pipeline, data
    )
    signature = infer_signature(
        data,
        generated_output,
        parameters,
    )

    mlflow.transformers.save_model(
        text2text_generation_pipeline,
        path=model_path,
        signature=signature,
    )
    pyfunc_loaded = mlflow.pyfunc.load_model(model_path)

    # parameters saved with ModelSignature is applied by default
    res = pyfunc_loaded.predict(data)
    res2 = pyfunc_loaded.predict(data, parameters)
    assert res == res2


def test_text2text_generation_pipeline_with_params_with_errors(
    text2text_generation_pipeline, model_path
):
    data = "muppet keyboard type"
    parameters = {"top_k": 2, "num_beams": 5, "invalid_param": "invalid_param", "do_sample": True}
    generated_output = mlflow.transformers.generate_signature_output(
        text2text_generation_pipeline, data
    )

    mlflow.transformers.save_model(
        text2text_generation_pipeline,
        path=model_path,
        signature=infer_signature(
            data,
            generated_output,
            parameters,
        ),
    )
    pyfunc_loaded = mlflow.pyfunc.load_model(model_path)
    with pytest.raises(
        MlflowException,
        match=r"The params provided to the `predict` method are "
        r"not valid for pipeline Text2TextGenerationPipeline.",
    ):
        pyfunc_loaded.predict(data, parameters)

    # Type validation of params failure
    with pytest.raises(MlflowException, match=r"Invalid parameters found"):
        pyfunc_loaded.predict(data, {"top_k": "2"})


def test_text2text_generation_pipeline_with_inferred_schema(text2text_generation_pipeline):
    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(text2text_generation_pipeline, name="my_model")
    pyfunc_loaded = mlflow.pyfunc.load_model(model_info.model_uri)

    assert pyfunc_loaded.predict("muppet board nails hammer")[0].startswith("A hammer")


@pytest.mark.parametrize(
    "invalid_data",
    [
        ({"answer": "something", "context": ["nothing", "that", "makes", "sense"]}),
        ([{"answer": ["42"], "context": "life"}, {"unmatched": "keys", "cause": "failure"}]),
    ],
)
def test_invalid_input_to_text2text_pipeline(text2text_generation_pipeline, invalid_data):
    # Adding this validation test due to the fact that we're constructing the input to the
    # Pipeline. The Pipeline requires a format of a pseudo-dict-like string. An example of
    # a valid input string: "answer: green. context: grass is primarily green in color."
    # We generate this string from a dict or generate a list of these strings from a list of
    # dictionaries.
    with pytest.raises(
        MlflowException, match=r"An invalid type has been supplied: .+\. Please supply"
    ):
        infer_signature(
            invalid_data,
            mlflow.transformers.generate_signature_output(
                text2text_generation_pipeline, invalid_data
            ),
        )


@pytest.mark.parametrize(
    "data", ["Generative models are", (["Generative models are", "Computers are"])]
)
def test_text_generation_pipeline(text_generation_pipeline, model_path, data):
    signature = infer_signature(
        data, mlflow.transformers.generate_signature_output(text_generation_pipeline, data)
    )

    model_config = {
        "prefix": "software",
        "top_k": 2,
        "num_beams": 5,
        "max_length": 30,
        "temperature": 0.62,
        "top_p": 0.85,
        "repetition_penalty": 1.15,
    }
    mlflow.transformers.save_model(
        text_generation_pipeline,
        path=model_path,
        model_config=model_config,
        signature=signature,
    )
    pyfunc_loaded = mlflow.pyfunc.load_model(model_path)

    inference = pyfunc_loaded.predict(data)

    if isinstance(data, list):
        assert inference[0].startswith(data[0])
        assert inference[1].startswith(data[1])
    else:
        assert inference[0].startswith(data)

    pd_input = pd.DataFrame([data], index=[0]) if isinstance(data, str) else pd.DataFrame(data)
    pd_inference = pyfunc_loaded.predict(pd_input)

    if isinstance(data, list):
        assert pd_inference[0].startswith(data[0])
        assert pd_inference[1].startswith(data[1])
    else:
        assert pd_inference[0].startswith(data)


@pytest.mark.parametrize(
    "invalid_data",
    [
        ({"my_input": "something to predict"}),
        ([{"bogus_input": "invalid"}, "not_valid"]),
        (["tell me a story", {"of": "a properly configured pipeline input"}]),
    ],
)
def test_invalid_input_to_text_generation_pipeline(text_generation_pipeline, invalid_data):
    if isinstance(invalid_data, list):
        match = "If supplying a list, all values must be of string type"
    else:
        match = "The input data is of an incorrect type"
    with pytest.raises(MlflowException, match=match):
        infer_signature(
            invalid_data,
            mlflow.transformers.generate_signature_output(text_generation_pipeline, invalid_data),
        )


@pytest.mark.parametrize(
    ("inference_payload", "result"),
    [
        ("Riding a <mask> on the beach is fun!", ["bike"]),
        (["If I had <mask>, I would fly to the top of a mountain"], ["wings"]),
        (
            ["I use stacks of <mask> to buy things", "I <mask> the whole bowl of cherries"],
            ["cash", "ate"],
        ),
    ],
)
def test_fill_mask_pipeline(fill_mask_pipeline, model_path, inference_payload, result):
    signature = infer_signature(
        inference_payload,
        mlflow.transformers.generate_signature_output(fill_mask_pipeline, inference_payload),
    )

    mlflow.transformers.save_model(fill_mask_pipeline, path=model_path, signature=signature)
    pyfunc_loaded = mlflow.pyfunc.load_model(model_path)

    inference = pyfunc_loaded.predict(inference_payload)
    assert inference == result

    if len(inference_payload) > 1 and isinstance(inference_payload, list):
        pd_input = pd.DataFrame([{"inputs": v} for v in inference_payload])
    elif isinstance(inference_payload, list) and len(inference_payload) == 1:
        pd_input = pd.DataFrame([{"inputs": v} for v in inference_payload], index=[0])
    else:
        pd_input = pd.DataFrame({"inputs": inference_payload}, index=[0])

    pd_inference = pyfunc_loaded.predict(pd_input)
    assert pd_inference == result


def test_fill_mask_pipeline_with_multiple_masks(fill_mask_pipeline, model_path):
    data = ["I <mask> the whole <mask> of <mask>", "I <mask> the whole <mask> of <mask>"]

    mlflow.transformers.save_model(fill_mask_pipeline, path=model_path)
    pyfunc_loaded = mlflow.pyfunc.load_model(model_path)

    inference = pyfunc_loaded.predict(data)
    assert len(inference) == 2
    assert all(len(value) == 3 for value in inference)


@pytest.mark.parametrize(
    "invalid_data",
    [
        ({"a": "b"}),
        ([{"a": "b"}, [{"a": "c"}]]),
    ],
)
def test_invalid_input_to_fill_mask_pipeline(fill_mask_pipeline, invalid_data):
    if isinstance(invalid_data, list):
        match = "Invalid data submission. Ensure all"
    else:
        match = "The input data is of an incorrect type"
    with pytest.raises(MlflowException, match=match):
        infer_signature(
            invalid_data,
            mlflow.transformers.generate_signature_output(fill_mask_pipeline, invalid_data),
        )


@pytest.mark.parametrize(
    "data",
    [
        {
            "sequences": "I love the latest update to this IDE!",
            "candidate_labels": ["happy", "sad"],
        },
        {
            "sequences": ["My dog loves to eat spaghetti", "My dog hates going to the vet"],
            "candidate_labels": ["happy", "sad"],
            "hypothesis_template": "This example talks about how the dog is {}",
        },
    ],
)
def test_zero_shot_classification_pipeline(zero_shot_pipeline, model_path, data):
    # NB: The list submission for this pipeline type can accept json-encoded lists or lists within
    # the values of the dictionary.
    signature = infer_signature(
        data, mlflow.transformers.generate_signature_output(zero_shot_pipeline, data)
    )

    mlflow.transformers.save_model(zero_shot_pipeline, model_path, signature=signature)

    loaded_pyfunc = mlflow.pyfunc.load_model(model_path)
    inference = loaded_pyfunc.predict(data)

    assert isinstance(inference, pd.DataFrame)
    if isinstance(data["sequences"], str):
        assert len(inference) == len(data["candidate_labels"])
    else:
        assert len(inference) == len(data["sequences"]) * len(data["candidate_labels"])


@pytest.mark.parametrize(
    "query",
    [
        "What should we order more of?",
        [
            "What is our highest sales?",
            "What should we order more of?",
        ],
    ],
)
def test_table_question_answering_pipeline(table_question_answering_pipeline, model_path, query):
    table = {
        "Fruit": ["Apples", "Bananas", "Oranges", "Watermelon", "Blueberries"],
        "Sales": ["1230945.55", "86453.12", "11459.23", "8341.23", "2325.88"],
        "Inventory": ["910", "4589", "11200", "80", "3459"],
    }
    json_table = json.dumps(table)
    data = {"query": query, "table": json_table}
    signature = infer_signature(
        data, mlflow.transformers.generate_signature_output(table_question_answering_pipeline, data)
    )

    mlflow.transformers.save_model(
        table_question_answering_pipeline, model_path, signature=signature
    )
    loaded = mlflow.pyfunc.load_model(model_path)

    inference = loaded.predict(data)
    assert len(inference) == 1 if isinstance(query, str) else len(query)

    pd_input = pd.DataFrame([data])
    pd_inference = loaded.predict(pd_input)
    assert pd_inference is not None


def test_custom_code_pipeline(custom_code_pipeline, model_path):
    data = "hello"

    signature = infer_signature(
        data, mlflow.transformers.generate_signature_output(custom_code_pipeline, data)
    )

    mlflow.transformers.save_model(
        custom_code_pipeline,
        path=model_path,
        signature=signature,
    )

    # just test that it doesn't blow up when performing inference
    pyfunc_loaded = mlflow.pyfunc.load_model(model_path)
    pyfunc_pred = pyfunc_loaded.predict(data)
    assert isinstance(pyfunc_pred[0][0], float)

    transformers_loaded = mlflow.transformers.load_model(model_path)
    transformers_pred = transformers_loaded(data)
    assert pyfunc_pred[0][0] == transformers_pred[0][0][0]


def test_custom_components_pipeline(custom_components_pipeline, model_path):
    data = "hello"

    signature = infer_signature(
        data, mlflow.transformers.generate_signature_output(custom_components_pipeline, data)
    )

    components = {
        "model": custom_components_pipeline.model,
        "tokenizer": custom_components_pipeline.tokenizer,
        "feature_extractor": custom_components_pipeline.feature_extractor,
    }

    # Pass the task explicitly rather than relying on inference from the Hub:
    # the backing model's `pipeline_tag` is now `image-feature-extraction`, so
    # task inference would build an unsupported pipeline type. The fixture
    # pipeline is already built with the `feature-extraction` task.
    mlflow.transformers.save_model(
        transformers_model=components,
        path=model_path,
        signature=signature,
        task=custom_components_pipeline.task,
    )

    pyfunc_loaded = mlflow.pyfunc.load_model(model_path)
    pyfunc_pred = pyfunc_loaded.predict(data)
    assert isinstance(pyfunc_pred[0][0], float)

    transformers_loaded = mlflow.transformers.load_model(model_path)
    transformers_pred = transformers_loaded(data)
    assert pyfunc_pred[0][0] == transformers_pred[0][0][0]

    # assert that all the reloaded components exist
    # and have the same class name as pre-save
    for name, component in components.items():
        assert component.__class__.__name__ == getattr(transformers_loaded, name).__class__.__name__


@pytest.mark.parametrize(
    ("data", "result"),
    [
        ("I've got a lovely bunch of coconuts!", ["Ich habe eine schöne Haufe von Kokos!"]),
        (
            [
                "I am the very model of a modern major general",
                "Once upon a time, there was a little turtle",
            ],
            [
                "Ich bin das Modell eines modernen Generals.",
                "Einmal gab es eine kleine Schildkröte.",
            ],
        ),
    ],
)
def test_translation_pipeline(translation_pipeline, model_path, data, result):
    signature = infer_signature(
        data, mlflow.transformers.generate_signature_output(translation_pipeline, data)
    )

    mlflow.transformers.save_model(translation_pipeline, path=model_path, signature=signature)
    pyfunc_loaded = mlflow.pyfunc.load_model(model_path)
    inference = pyfunc_loaded.predict(data)
    assert inference == result

    if len(data) > 1 and isinstance(data, list):
        pd_input = pd.DataFrame([{"inputs": v} for v in data])
    elif isinstance(data, list) and len(data) == 1:
        pd_input = pd.DataFrame([{"inputs": v} for v in data], index=[0])
    else:
        pd_input = pd.DataFrame({"inputs": data}, index=[0])

    pd_inference = pyfunc_loaded.predict(pd_input)
    assert pd_inference == result


@pytest.mark.parametrize(
    "data",
    [
        "There once was a boy",
        ["Dolly isn't just a sheep anymore"],
        ["Baking cookies is quite easy", "Writing unittests is good for"],
    ],
)
def test_summarization_pipeline(summarizer_pipeline, model_path, data):
    model_config = {
        "top_k": 2,
        "num_beams": 5,
        "max_length": 90,
        "temperature": 0.62,
        "top_p": 0.85,
        "repetition_penalty": 1.15,
    }
    signature = infer_signature(
        data, mlflow.transformers.generate_signature_output(summarizer_pipeline, data)
    )

    mlflow.transformers.save_model(
        summarizer_pipeline, path=model_path, model_config=model_config, signature=signature
    )
    pyfunc_loaded = mlflow.pyfunc.load_model(model_path)

    inference = pyfunc_loaded.predict(data)
    if isinstance(data, list) and len(data) > 1:
        for i, entry in enumerate(data):
            assert inference[i].strip().startswith(entry)
    elif isinstance(data, list) and len(data) == 1:
        assert inference[0].strip().startswith(data[0])
    else:
        assert inference[0].strip().startswith(data)

    if len(data) > 1 and isinstance(data, list):
        pd_input = pd.DataFrame([{"inputs": v} for v in data])
    elif isinstance(data, list) and len(data) == 1:
        pd_input = pd.DataFrame([{"inputs": v} for v in data], index=[0])
    else:
        pd_input = pd.DataFrame({"inputs": data}, index=[0])

    pd_inference = pyfunc_loaded.predict(pd_input)
    if isinstance(data, list) and len(data) > 1:
        for i, entry in enumerate(data):
            assert pd_inference[i].strip().startswith(entry)
    elif isinstance(data, list) and len(data) == 1:
        assert pd_inference[0].strip().startswith(data[0])
    else:
        assert pd_inference[0].strip().startswith(data)


@pytest.mark.parametrize(
    "data",
    [
        "I'm telling you that Han shot first!",
        [
            "I think this sushi might have gone off",
            "That gym smells like feet, hot garbage, and sadness",
            "I love that we have a moon",
        ],
        [{"text": "test1", "text_pair": "test2"}],
        [{"text": "test1", "text_pair": "pair1"}, {"text": "test2", "text_pair": "pair2"}],
    ],
)
def test_classifier_pipeline(text_classification_pipeline, model_path, data):
    signature = infer_signature(
        data, mlflow.transformers.generate_signature_output(text_classification_pipeline, data)
    )
    mlflow.transformers.save_model(
        text_classification_pipeline, path=model_path, signature=signature
    )

    pyfunc_loaded = mlflow.pyfunc.load_model(model_path)
    inference = pyfunc_loaded.predict(data)

    # verify that native transformers outputs match the pyfunc return values
    native_inference = text_classification_pipeline(data)
    inference_dict = inference.to_dict()

    if isinstance(data, str):
        assert len(inference) == 1
        assert inference_dict["label"][0] == native_inference[0]["label"]
        assert inference_dict["score"][0] == native_inference[0]["score"]
    else:
        assert len(inference) == len(data)
        for key in ["score", "label"]:
            for value in range(0, len(data)):
                if key == "label":
                    assert inference_dict[key][value] == native_inference[value][key]
                else:
                    assert math.isclose(
                        native_inference[value][key], inference_dict[key][value], rel_tol=1e-6
                    )


@pytest.mark.parametrize(
    ("data", "result"),
    [
        (
            "I have a dog and his name is Willy!",
            ["PRON,VERB,DET,NOUN,CCONJ,PRON,NOUN,AUX,PROPN,PUNCT"],
        ),
        (["I like turtles"], ["PRON,VERB,NOUN"]),
        (
            ["We are the knights who say nee!", "Houston, we may have a problem."],
            [
                "PRON,AUX,DET,PROPN,PRON,VERB,INTJ,PUNCT",
                "PROPN,PUNCT,PRON,AUX,VERB,DET,NOUN,PUNCT",
            ],
        ),
    ],
)
@pytest.mark.parametrize("pipeline_name", ["ner_pipeline", "ner_pipeline_aggregation"])
def test_ner_pipeline(pipeline_name, model_path, data, result, request):
    pipeline = request.getfixturevalue(pipeline_name)

    signature = infer_signature(data, mlflow.transformers.generate_signature_output(pipeline, data))

    mlflow.transformers.save_model(pipeline, model_path, signature=signature)
    loaded_pyfunc = mlflow.pyfunc.load_model(model_path)
    inference = loaded_pyfunc.predict(data)

    assert inference == result

    if len(data) > 1 and isinstance(data, list):
        pd_input = pd.DataFrame([{"inputs": v} for v in data])
    elif isinstance(data, list) and len(data) == 1:
        pd_input = pd.DataFrame([{"inputs": v} for v in data], index=[0])
    else:
        pd_input = pd.DataFrame({"inputs": data}, index=[0])
    pd_inference = loaded_pyfunc.predict(pd_input)
    assert pd_inference == result


@pytest.mark.skipif(
    _try_import_conversational_pipeline() is None,
    reason="Conversation model is deprecated and removed.",
)
def test_conversational_pipeline(conversational_pipeline, model_path):
    assert mlflow.transformers._is_conversational_pipeline(conversational_pipeline)

    signature = infer_signature(
        "Hi there!",
        mlflow.transformers.generate_signature_output(conversational_pipeline, "Hi there!"),
    )

    mlflow.transformers.save_model(conversational_pipeline, model_path, signature=signature)
    loaded_pyfunc = mlflow.pyfunc.load_model(model_path)

    first_response = loaded_pyfunc.predict("What is the best way to get to Antarctica?")

    assert first_response == "The best way would be to go to space."

    second_response = loaded_pyfunc.predict("What kind of boat should I use?")

    assert second_response == "The best way to get to space would be to reach out and touch it."

    # Test that a new loaded instance has no context.
    loaded_again_pyfunc = mlflow.pyfunc.load_model(model_path)
    third_response = loaded_again_pyfunc.predict("What kind of boat should I use?")

    assert third_response == "The one with the guns."

    fourth_response = loaded_again_pyfunc.predict("Can I use it to go to the moon?")

    assert fourth_response == "Sure."


def test_qa_pipeline_pyfunc_predict(small_qa_pipeline):
    artifact_path = "qa_model"
    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            small_qa_pipeline,
            name=artifact_path,
        )

    inference_payload = json.dumps({
        "inputs": {
            "question": [
                "What color is it?",
                "How do the people go?",
                "What does the 'wolf' howl at?",
            ],
            "context": [
                "Some people said it was green but I know that it's pink.",
                "The people on the bus go up and down. Up and down.",
                "The pack of 'wolves' stood on the cliff and a 'lone wolf' howled at "
                "the moon for hours.",
            ],
        }
    })
    response = pyfunc_serve_and_score_model(
        model_info.model_uri,
        data=inference_payload,
        content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON,
        extra_args=["--env-manager", "local"],
    )
    values = PredictionsResponse.from_json(response.content.decode("utf-8")).get_predictions()

    assert values.to_dict(orient="records") == [{0: "pink"}, {0: "up and down"}, {0: "the moon"}]

    inference_payload = json.dumps({
        "inputs": {
            "question": "Who's house?",
            "context": "The house is owned by a man named Run.",
        }
    })

    response = pyfunc_serve_and_score_model(
        model_info.model_uri,
        data=inference_payload,
        content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON,
        extra_args=["--env-manager", "local"],
    )
    values = PredictionsResponse.from_json(response.content.decode("utf-8")).get_predictions()

    assert values.to_dict(orient="records") == [{0: "Run"}]


@pytest.mark.parametrize(
    ("input_image", "result"),
    [
        (str(image_file_path), False),
        (image_url, False),
        ("base64", True),
        ("random string", False),
    ],
)
def test_vision_is_base64_image(input_image, result):
    if input_image == "base64":
        input_image = base64.b64encode(image_file_path.read_bytes()).decode("utf-8")
    assert _TransformersWrapper.is_base64_image(input_image) == result


@pytest.mark.parametrize(
    "inference_payload",
    [
        [str(image_file_path)],
        [image_url],
        "base64",
        pytest.param(
            "base64_encodebytes",
            marks=pytest.mark.skipif(
                Version(transformers.__version__) < Version("4.41"),
                reason="base64 encodebytes feature not present",
            ),
        ),
    ],
)
def test_vision_pipeline_pyfunc_predict(small_vision_model, inference_payload):
    if inference_payload == "base64":
        inference_payload = [
            base64.b64encode(image_file_path.read_bytes()).decode("utf-8"),
        ]
    elif inference_payload == "base64_encodebytes":
        inference_payload = [
            base64.encodebytes(image_file_path.read_bytes()).decode("utf-8"),
        ]
    artifact_path = "image_classification_model"

    # Log the image classification model
    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            small_vision_model,
            name=artifact_path,
        )
    pyfunc_inference_payload = json.dumps({"inputs": inference_payload})
    response = pyfunc_serve_and_score_model(
        model_info.model_uri,
        data=pyfunc_inference_payload,
        content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON,
        extra_args=["--env-manager", "local"],
    )

    predictions = PredictionsResponse.from_json(response.content.decode("utf-8")).get_predictions()

    transformers_loaded_model = mlflow.transformers.load_model(model_info.model_uri)
    expected_predictions = transformers_loaded_model.predict(inference_payload)

    assert [list(pred.values()) for pred in predictions.to_dict("records")] == expected_predictions


def test_classifier_pipeline_pyfunc_predict(text_classification_pipeline):
    artifact_path = "text_classifier_model"
    data = [
        "I think this sushi might have gone off",
        "That gym smells like feet, hot garbage, and sadness",
        "I love that we have a moon",
        "I 'love' debugging subprocesses",
        'Quote "in" the string',
    ]
    signature = infer_signature(data)
    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            text_classification_pipeline,
            name=artifact_path,
            signature=signature,
        )

    response = pyfunc_serve_and_score_model(
        model_info.model_uri,
        data=json.dumps({"inputs": data}),
        content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON,
        extra_args=["--env-manager", "local"],
    )
    values = PredictionsResponse.from_json(response.content.decode("utf-8")).get_predictions()

    assert len(values.to_dict()) == 2
    assert len(values.to_dict()["score"]) == 5

    # test simple string input
    inference_payload = json.dumps({"inputs": ["testing"]})

    response = pyfunc_serve_and_score_model(
        model_info.model_uri,
        data=inference_payload,
        content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON,
        extra_args=["--env-manager", "local"],
    )
    values = PredictionsResponse.from_json(response.content.decode("utf-8")).get_predictions()

    assert len(values.to_dict()) == 2
    assert len(values.to_dict()["score"]) == 1

    # Test the alternate TextClassificationPipeline input structure where text_pair is used
    # and ensure that model serving and direct native inference match
    inference_data = [
        {"text": "test1", "text_pair": "pair1"},
        {"text": "test2", "text_pair": "pair2"},
        {"text": "test 'quote", "text_pair": "pair 'quote'"},
    ]
    signature = infer_signature(inference_data)
    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            text_classification_pipeline,
            name=artifact_path,
            signature=signature,
        )

    inference_payload = json.dumps({"inputs": inference_data})
    response = pyfunc_serve_and_score_model(
        model_info.model_uri,
        data=inference_payload,
        content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON,
        extra_args=["--env-manager", "local"],
    )
    values = PredictionsResponse.from_json(response.content.decode("utf-8")).get_predictions()
    values_dict = values.to_dict()
    native_predict = text_classification_pipeline(inference_data)

    # validate that the pyfunc served model registers text_pair in the same manner as native
    for key in ["score", "label"]:
        for value in [0, 1]:
            if key == "label":
                assert values_dict[key][value] == native_predict[value][key]
            else:
                assert math.isclose(
                    values_dict[key][value], native_predict[value][key], rel_tol=1e-6
                )


def test_zero_shot_pipeline_pyfunc_predict(zero_shot_pipeline):
    artifact_path = "zero_shot_classifier_model"
    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            zero_shot_pipeline,
            name=artifact_path,
        )
        model_uri = model_info.model_uri

    inference_payload = json.dumps({
        "inputs": {
            "sequences": "My dog loves running through troughs of spaghetti with his mouth open",
            "candidate_labels": ["happy", "sad"],
            "hypothesis_template": "This example talks about how the dog is {}",
        }
    })

    response = pyfunc_serve_and_score_model(
        model_uri,
        data=inference_payload,
        content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON,
        extra_args=["--env-manager", "local"],
    )
    values = PredictionsResponse.from_json(response.content.decode("utf-8")).get_predictions()

    assert len(values.to_dict()) == 3
    assert len(values.to_dict()["labels"]) == 2

    inference_payload = json.dumps({
        "inputs": {
            "sequences": [
                "My dog loves to eat spaghetti",
                "My dog hates going to the vet",
                "My 'hamster' loves to play with my 'friendly' dog",
            ],
            "candidate_labels": '["happy", "sad"]',
            "hypothesis_template": "This example talks about how the dog is {}",
        }
    })
    response = pyfunc_serve_and_score_model(
        model_uri,
        data=inference_payload,
        content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON,
        extra_args=["--env-manager", "local"],
    )
    values = PredictionsResponse.from_json(response.content.decode("utf-8")).get_predictions()

    assert len(values.to_dict()) == 3
    assert len(values.to_dict()["labels"]) == 6


def test_table_question_answering_pyfunc_predict(table_question_answering_pipeline):
    artifact_path = "table_qa_model"
    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            table_question_answering_pipeline,
            name=artifact_path,
        )

    table = {
        "Fruit": ["Apples", "Bananas", "Oranges", "Watermelon 'small'", "Blueberries"],
        "Sales": ["1230945.55", "86453.12", "11459.23", "8341.23", "2325.88"],
        "Inventory": ["910", "4589", "11200", "80", "3459"],
    }

    inference_payload = json.dumps({
        "inputs": {
            "query": "What should we order more of?",
            "table": table,
        }
    })

    response = pyfunc_serve_and_score_model(
        model_info.model_uri,
        data=inference_payload,
        content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON,
        extra_args=["--env-manager", "local"],
    )
    values = PredictionsResponse.from_json(response.content.decode("utf-8")).get_predictions()

    assert len(values.to_dict(orient="records")) == 1

    inference_payload = json.dumps({
        "inputs": {
            "query": [
                "What is our highest sales?",
                "What should we order more of?",
                "Which 'fruit' has the 'highest' 'sales'?",
            ],
            "table": table,
        }
    })
    response = pyfunc_serve_and_score_model(
        model_info.model_uri,
        data=inference_payload,
        content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON,
        extra_args=["--env-manager", "local"],
    )
    values = PredictionsResponse.from_json(response.content.decode("utf-8")).get_predictions()

    assert len(values.to_dict(orient="records")) == 3


def test_feature_extraction_pipeline(feature_extraction_pipeline):
    sentences = ["hi", "hello"]
    signature = infer_signature(
        sentences,
        mlflow.transformers.generate_signature_output(feature_extraction_pipeline, sentences),
    )

    artifact_path = "feature_extraction_pipeline"
    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            feature_extraction_pipeline,
            name=artifact_path,
            signature=signature,
            input_example=["A sentence", "Another sentence"],
        )

    # Load as native
    loaded_pipeline = mlflow.transformers.load_model(model_info.model_uri)

    inference_single = "Testing"
    inference_mult = ["Testing something", "Testing something else"]

    pred = loaded_pipeline(inference_single)
    assert len(pred[0][0]) > 10
    assert isinstance(pred[0][0][0], float)

    pred_multiple = loaded_pipeline(inference_mult)
    assert len(pred_multiple[0][0]) > 2
    assert isinstance(pred_multiple[0][0][0][0], float)

    loaded_pyfunc = mlflow.pyfunc.load_model(model_info.model_uri)

    pyfunc_pred = loaded_pyfunc.predict(inference_single)

    assert isinstance(pyfunc_pred, np.ndarray)

    assert np.array_equal(np.array(pred[0]), pyfunc_pred)

    pyfunc_pred_multiple = loaded_pyfunc.predict(inference_mult)

    assert np.array_equal(np.array(pred_multiple[0][0]), pyfunc_pred_multiple)


def test_feature_extraction_pipeline_pyfunc_predict(feature_extraction_pipeline):
    artifact_path = "feature_extraction"
    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            feature_extraction_pipeline,
            name=artifact_path,
        )

    inference_payload = json.dumps({"inputs": ["sentence one", "sentence two"]})

    response = pyfunc_serve_and_score_model(
        model_info.model_uri,
        data=inference_payload,
        content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON,
        extra_args=["--env-manager", "local"],
    )
    values = PredictionsResponse.from_json(response.content.decode("utf-8")).get_predictions()

    assert len(values.columns) == 384
    assert len(values) == 4

    inference_payload = json.dumps({"inputs": "sentence three"})

    response = pyfunc_serve_and_score_model(
        model_info.model_uri,
        data=inference_payload,
        content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON,
        extra_args=["--env-manager", "local"],
    )
    assert response.status_code == 200
    prediction = PredictionsResponse.from_json(response.content.decode("utf-8")).get_predictions()
    assert len(prediction.columns) == 384
    assert len(prediction) == 4


def test_loading_unsupported_pipeline_type_as_pyfunc(small_multi_modal_pipeline, model_path):
    mlflow.transformers.save_model(small_multi_modal_pipeline, model_path)
    with pytest.raises(MlflowException, match='Model does not have the "python_function" flavor'):
        mlflow.pyfunc.load_model(model_path)


def test_pyfunc_input_validations(mock_pyfunc_wrapper):
    def ensure_raises(data, match):
        with pytest.raises(MlflowException, match=match):
            mock_pyfunc_wrapper._validate_str_or_list_str(data)

    match1 = "The input data is of an incorrect type"
    match2 = "If supplying a list, all values must"
    ensure_raises({"a": "b"}, match1)
    ensure_raises(("a", "b"), match1)
    ensure_raises({"a", "b"}, match1)
    ensure_raises(True, match1)
    ensure_raises(12, match1)
    ensure_raises([1, 2, 3], match2)
    ensure_raises([{"a", "b"}], match2)
    ensure_raises([["a", "b", "c'"]], match2)
    ensure_raises([{"a": "b"}, {"a": "c"}], match2)
    ensure_raises([[1], [2]], match2)


def test_pyfunc_json_encoded_dict_parsing(mock_pyfunc_wrapper):
    plain_dict = {"a": "b", "b": "c"}
    list_dict = [plain_dict, plain_dict]

    plain_input = {"in": json.dumps(plain_dict)}
    list_input = {"in": json.dumps(list_dict)}

    plain_parsed = mock_pyfunc_wrapper._parse_json_encoded_dict_payload_to_dict(plain_input, "in")
    assert plain_parsed == {"in": plain_dict}

    list_parsed = mock_pyfunc_wrapper._parse_json_encoded_dict_payload_to_dict(list_input, "in")
    assert list_parsed == {"in": list_dict}

    invalid_parsed = mock_pyfunc_wrapper._parse_json_encoded_dict_payload_to_dict(
        plain_input, "invalid"
    )
    assert invalid_parsed != {"in": plain_dict}
    assert invalid_parsed == plain_input


def test_pyfunc_json_encoded_list_parsing(mock_pyfunc_wrapper):
    plain_list = ["a", "b", "c"]
    nested_list = [plain_list, plain_list]
    list_dict = [{"a": "b"}, {"a": "c"}]

    plain_input = {"in": json.dumps(plain_list)}
    nested_input = {"in": json.dumps(nested_list)}
    list_dict_input = {"in": json.dumps(list_dict)}

    plain_parsed = mock_pyfunc_wrapper._parse_json_encoded_list(plain_input, "in")
    assert plain_parsed == {"in": plain_list}

    nested_parsed = mock_pyfunc_wrapper._parse_json_encoded_list(nested_input, "in")
    assert nested_parsed == {"in": nested_list}

    list_dict_parsed = mock_pyfunc_wrapper._parse_json_encoded_list(list_dict_input, "in")
    assert list_dict_parsed == {"in": list_dict}

    with pytest.raises(MlflowException, match="Invalid key in inference payload. The "):
        mock_pyfunc_wrapper._parse_json_encoded_list(list_dict_input, "invalid")


def test_pyfunc_text_to_text_input(mock_pyfunc_wrapper):
    text2text_input = {"context": "a", "answer": "b"}
    parsed_input = mock_pyfunc_wrapper._parse_text2text_input(text2text_input)
    assert parsed_input == "context: a answer: b"

    text2text_input_list = [text2text_input, text2text_input]
    parsed_input_list = mock_pyfunc_wrapper._parse_text2text_input(text2text_input_list)
    assert parsed_input_list == ["context: a answer: b", "context: a answer: b"]

    parsed_with_inputs = mock_pyfunc_wrapper._parse_text2text_input({"inputs": "a"})
    assert parsed_with_inputs == ["a"]

    parsed_str = mock_pyfunc_wrapper._parse_text2text_input("a")
    assert parsed_str == "a"

    parsed_list_str = mock_pyfunc_wrapper._parse_text2text_input(["a", "b"])
    assert parsed_list_str == ["a", "b"]

    with pytest.raises(MlflowException, match="An invalid type has been supplied"):
        mock_pyfunc_wrapper._parse_text2text_input([1, 2, 3])

    with pytest.raises(MlflowException, match="An invalid type has been supplied"):
        mock_pyfunc_wrapper._parse_text2text_input([{"a": [{"b": "c"}]}])


def test_pyfunc_qa_input(mock_pyfunc_wrapper):
    single_input = {"question": "a", "context": "b"}
    parsed_single_input = mock_pyfunc_wrapper._parse_question_answer_input(single_input)
    assert parsed_single_input == single_input

    multi_input = [single_input, single_input]
    parsed_multi_input = mock_pyfunc_wrapper._parse_question_answer_input(multi_input)
    assert parsed_multi_input == multi_input

    with pytest.raises(MlflowException, match="Invalid keys were submitted. Keys must"):
        mock_pyfunc_wrapper._parse_question_answer_input({"q": "a", "c": "b"})

    with pytest.raises(MlflowException, match="An invalid type has been supplied"):
        mock_pyfunc_wrapper._parse_question_answer_input("a")

    with pytest.raises(MlflowException, match="An invalid type has been supplied"):
        mock_pyfunc_wrapper._parse_question_answer_input(["a", "b", "c"])


def test_list_of_dict_to_list_of_str_parsing(mock_pyfunc_wrapper):
    # Test with a single list of dictionaries
    output_data = [{"a": "foo"}, {"a": "bar"}, {"a": "baz"}]
    expected_output = ["foo", "bar", "baz"]
    assert (
        mock_pyfunc_wrapper._parse_lists_of_dict_to_list_of_str(output_data, "a") == expected_output
    )

    # Test with a nested list of dictionaries
    output_data = [
        {"a": "foo", "b": [{"a": "bar"}]},
        {"a": "baz", "b": [{"a": "qux"}]},
    ]
    expected_output = ["foo", "bar", "baz", "qux"]
    assert (
        mock_pyfunc_wrapper._parse_lists_of_dict_to_list_of_str(output_data, "a") == expected_output
    )

    # Test with nested list with exclusion data
    output_data = [
        {"a": "valid", "b": [{"a": "another valid"}, {"b": "invalid"}]},
        {"a": "valid 2", "b": [{"a": "another valid 2"}, {"c": "invalid"}]},
    ]
    expected_output = ["valid", "another valid", "valid 2", "another valid 2"]
    assert (
        mock_pyfunc_wrapper._parse_lists_of_dict_to_list_of_str(output_data, "a") == expected_output
    )


def test_parsing_tokenizer_output(mock_pyfunc_wrapper):
    output_data = [{"a": "b"}, {"a": "c"}, {"a": "d"}]
    expected_output = "b,c,d"
    assert mock_pyfunc_wrapper._parse_tokenizer_output(output_data, {"a"}) == expected_output

    output_data = [output_data, output_data]
    expected_output = [expected_output, expected_output]
    assert mock_pyfunc_wrapper._parse_tokenizer_output(output_data, {"a"}) == expected_output


def test_parse_list_of_multiple_dicts(mock_pyfunc_wrapper):
    output_data = [{"a": "b", "d": "f"}, {"a": "z", "d": "g"}]
    target_dict_key = "a"
    expected_output = ["b"]

    assert (
        mock_pyfunc_wrapper._parse_list_of_multiple_dicts(output_data, target_dict_key)
        == expected_output
    )

    output_data = [
        [{"a": "c", "d": "q"}, {"a": "o", "d": "q"}, {"a": "d", "d": "q"}, {"a": "e", "d": "r"}],
        [{"a": "m", "d": "s"}, {"a": "e", "d": "t"}],
    ]
    target_dict_key = "a"
    expected_output = ["c", "m"]

    assert (
        mock_pyfunc_wrapper._parse_list_of_multiple_dicts(output_data, target_dict_key)
        == expected_output
    )


@pytest.mark.parametrize(
    (
        "pipeline_input",
        "pipeline_output",
        "expected_output",
        "flavor_config",
        "include_prompt",
        "collapse_whitespace",
    ),
    [
        (
            "What answers?",
            [{"generated_text": "What answers?\n\nA collection of\n\nanswers"}],
            "A collection of\n\nanswers",
            {"instance_type": "InstructionTextGenerationPipeline"},
            False,
            False,
        ),
        (
            "What answers?",
            [{"generated_text": "What answers?\n\nA collection of\n\nanswers"}],
            "A collection of answers",
            {"instance_type": "InstructionTextGenerationPipeline"},
            False,
            True,
        ),
        (
            "Hello!",
            [{"generated_text": "Hello!\n\nHow are you?"}],
            "How are you?",
            {"instance_type": "InstructionTextGenerationPipeline"},
            False,
            False,
        ),
        (
            "Hello!",
            [{"generated_text": "Hello!\n\nA: How are you?\n\n"}],
            "How are you?",
            {"instance_type": "InstructionTextGenerationPipeline"},
            False,
            True,
        ),
        (
            "Hello!",
            [{"generated_text": "Hello!\n\nA: How are you?\n\n"}],
            "Hello! A: How are you?",
            {"instance_type": "InstructionTextGenerationPipeline"},
            True,
            True,
        ),
        (
            "Hello!",
            [{"generated_text": "Hello!\n\nA: How\nare\nyou?\n\n"}],
            "How\nare\nyou?\n\n",
            {"instance_type": "InstructionTextGenerationPipeline"},
            False,
            False,
        ),
        (
            ["Hi!", "What's up?"],
            [[{"generated_text": "Hi!\n\nHello there"}, {"generated_text": "Not much, and you?"}]],
            ["Hello there", "Not much, and you?"],
            {"instance_type": "InstructionTextGenerationPipeline"},
            False,
            False,
        ),
        # Tests disabling parsing of newline characters
        (
            ["Hi!", "What's up?"],
            [
                [
                    {"generated_text": "Hi!\n\nHello there"},
                    {"generated_text": "What's up?\n\nNot much, and you?"},
                ]
            ],
            ["Hi!\n\nHello there", "What's up?\n\nNot much, and you?"],
            {"instance_type": "InstructionTextGenerationPipeline"},
            True,
            False,
        ),
        (
            "Hello!",
            [{"generated_text": "Hello!\n\nHow are you?"}],
            "Hello!\n\nHow are you?",
            {"instance_type": "InstructionTextGenerationPipeline"},
            True,
            False,
        ),
        # Tests a standard TextGenerationPipeline output
        (
            ["We like to", "Open the"],
            [
                [
                    {"generated_text": "We like to party"},
                    {"generated_text": "Open the door get on the floor everybody do the dinosaur"},
                ]
            ],
            ["We like to party", "Open the door get on the floor everybody do the dinosaur"],
            {"instance_type": "TextGenerationPipeline"},
            True,
            True,
        ),
        # Tests a standard TextGenerationPipeline output with setting "include_prompt" (noop)
        (
            ["We like to", "Open the"],
            [
                [
                    {"generated_text": "We like to party"},
                    {"generated_text": "Open the door get on the floor everybody do the dinosaur"},
                ]
            ],
            ["We like to party", "Open the door get on the floor everybody do the dinosaur"],
            {"instance_type": "TextGenerationPipeline"},
            False,
            False,
        ),
        # Test TextGenerationPipeline removes whitespace
        (
            ["We like to", "Open the"],
            [
                [
                    {"generated_text": "  We like   to    party"},
                    {
                        "generated_text": "Open the   door get on the floor   everybody    "
                        "do\nthe dinosaur"
                    },
                ]
            ],
            ["We like to party", "Open the door get on the floor everybody do the dinosaur"],
            {"instance_type": "TextGenerationPipeline"},
            False,
            True,
        ),
    ],
)
def test_parse_input_from_instruction_pipeline(
    mock_pyfunc_wrapper,
    pipeline_input,
    pipeline_output,
    expected_output,
    flavor_config,
    include_prompt,
    collapse_whitespace,
):
    assert (
        mock_pyfunc_wrapper._strip_input_from_response_in_instruction_pipelines(
            pipeline_input,
            pipeline_output,
            "generated_text",
            flavor_config,
            include_prompt,
            collapse_whitespace,
        )
        == expected_output
    )


@pytest.mark.parametrize(
    "flavor_config",
    [
        {"instance_type": "InstructionTextGenerationPipeline"},
        {"instance_type": "TextGenerationPipeline"},
    ],
)
def test_invalid_instruction_pipeline_parsing(mock_pyfunc_wrapper, flavor_config):
    prompt = "What is your favorite boba flavor?"

    bad_output = {"generated_text": ["Strawberry Milk Cap", "Honeydew with boba"]}

    with pytest.raises(MlflowException, match="Unable to parse the pipeline output. Expected"):
        mock_pyfunc_wrapper._strip_input_from_response_in_instruction_pipelines(
            prompt, bad_output, "generated_text", flavor_config, True
        )


@pytest.mark.skipif(RUNNING_IN_GITHUB_ACTIONS, reason=GITHUB_ACTIONS_SKIP_REASON)
def test_instructional_pipeline_no_prompt_in_output(model_path):
    architecture = "databricks/dolly-v2-3b"
    dolly = transformers.pipeline(model=architecture, trust_remote_code=True)

    mlflow.transformers.save_model(
        transformers_model=dolly,
        path=model_path,
        # Validate removal of prompt but inclusion of newlines by default
        model_config={"max_length": 100, "include_prompt": False},
        input_example="Hello, Dolly!",
    )

    pyfunc_loaded = mlflow.pyfunc.load_model(model_path)

    inference = pyfunc_loaded.predict("What is MLflow?")

    assert not inference[0].startswith("What is MLflow?")
    assert "\n" in inference[0]


@pytest.mark.skipif(RUNNING_IN_GITHUB_ACTIONS, reason=GITHUB_ACTIONS_SKIP_REASON)
def test_instructional_pipeline_no_prompt_in_output_and_removal_of_newlines(model_path):
    architecture = "databricks/dolly-v2-3b"
    dolly = transformers.pipeline(model=architecture, trust_remote_code=True)

    mlflow.transformers.save_model(
        transformers_model=dolly,
        path=model_path,
        # Validate removal of prompt but inclusion of newlines by default
        model_config={"max_length": 100, "include_prompt": False, "collapse_whitespace": True},
        input_example="Hello, Dolly!",
    )

    pyfunc_loaded = mlflow.pyfunc.load_model(model_path)

    inference = pyfunc_loaded.predict("What is MLflow?")

    assert not inference[0].startswith("What is MLflow?")
    assert "\n" not in inference[0]


@pytest.mark.skipif(RUNNING_IN_GITHUB_ACTIONS, reason=GITHUB_ACTIONS_SKIP_REASON)
def test_instructional_pipeline_with_prompt_in_output(model_path):
    architecture = "databricks/dolly-v2-3b"
    dolly = transformers.pipeline(model=architecture, trust_remote_code=True)

    mlflow.transformers.save_model(
        transformers_model=dolly,
        path=model_path,
        # test default propagation of `include_prompt`=True and `collapse_whitespace`=False
        model_config={"max_length": 100},
        input_example="Hello, Dolly!",
    )

    pyfunc_loaded = mlflow.pyfunc.load_model(model_path)

    inference = pyfunc_loaded.predict("What is MLflow?")

    assert inference[0].startswith("What is MLflow?")
    assert "\n\n" in inference[0]


def read_audio_data(format: str):
    datasets_path = pathlib.Path(__file__).resolve().parent.parent.joinpath("datasets")
    wav_file_path = datasets_path.joinpath("apollo11_launch.wav")
    if format == "float":
        audio, _ = librosa.load(wav_file_path, sr=16000)
        return audio
    elif format == "bytes":
        return wav_file_path.read_bytes()
    elif format == "file":
        return wav_file_path.as_posix()
    else:
        raise ValueError(f"Invalid format: {format}")


@pytest.mark.parametrize("input_format", ["float", "bytes", "file"])
@pytest.mark.parametrize("with_input_example", [True, False])
def test_whisper_model_predict(model_path, whisper_pipeline, input_format, with_input_example):
    if input_format == "float" and not with_input_example:
        pytest.skip("If the input format is float, the default signature must be overridden.")

    audio = read_audio_data(input_format)
    mlflow.transformers.save_model(
        transformers_model=whisper_pipeline,
        path=model_path,
        input_example=audio if with_input_example else None,
        save_pretrained=False,
    )

    # 1. Single prediction with native transformer pipeline
    loaded_pipeline = mlflow.transformers.load_model(model_path)
    transcription = loaded_pipeline(audio)
    assert transcription["text"].startswith(" 30")
    # strip the leading space
    expected_text = transcription["text"].lstrip()

    # 2. Single prediction with Pyfunc
    loaded_pyfunc = mlflow.pyfunc.load_model(model_path)
    pyfunc_transcription = loaded_pyfunc.predict(audio)[0]
    assert pyfunc_transcription == expected_text

    # 3. Batch prediction with Pyfunc. Float input format is not supported for batch prediction,
    # because our signature framework doesn't support a list of numpy array.
    if input_format != "float":
        batch_transcription = loaded_pyfunc.predict([audio, audio])
        assert len(batch_transcription) == 2
        assert all(ts == expected_text for ts in batch_transcription)


def test_whisper_model_serve_and_score(whisper_pipeline):
    # Request payload to the model serving endpoint contains base64 encoded audio data
    audio = read_audio_data("bytes")
    encoded_audio = base64.b64encode(audio).decode("ascii")

    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            whisper_pipeline,
            name="whisper",
            save_pretrained=False,
        )

    def _assert_response(response, length=1):
        preds = json.loads(response.content.decode("utf-8"))["predictions"]
        assert len(preds) == length
        assert all(pred.startswith("30") for pred in preds)

    with pyfunc_scoring_endpoint(
        model_info.model_uri,
        extra_args=["--env-manager", "local"],
    ) as endpoint:
        content_type = pyfunc_scoring_server.CONTENT_TYPE_JSON

        # Test payload with "inputs" key
        inputs_dict = {"inputs": [encoded_audio]}
        payload = json.dumps(inputs_dict)
        response = endpoint.invoke(payload, content_type=content_type)
        _assert_response(response)

        # Test payload with "dataframe_split" key
        inference_df = pd.DataFrame(pd.Series([encoded_audio], name="audio_file"))
        split_dict = {"dataframe_split": inference_df.to_dict(orient="split")}
        payload = json.dumps(split_dict)
        response = endpoint.invoke(payload, content_type=content_type)
        _assert_response(response)

        # Test payload with "dataframe_records" key
        records_dict = {"dataframe_records": inference_df.to_dict(orient="records")}
        payload = json.dumps(records_dict)
        response = endpoint.invoke(payload, content_type=content_type)
        _assert_response(response)

        # Test batch prediction
        inputs_dict = {"inputs": [encoded_audio, encoded_audio]}
        payload = json.dumps(inputs_dict)
        response = endpoint.invoke(payload, content_type=content_type)
        _assert_response(response, length=2)

        # Scoring with audio file URI is not supported yet (pyfunc prediction works tho)
        inputs_dict = {"inputs": [read_audio_data("file")]}
        payload = json.dumps(inputs_dict)
        response = endpoint.invoke(payload, content_type=content_type)
        response = json.loads(response.content.decode("utf-8"))
        assert response["error_code"] == "INVALID_PARAMETER_VALUE"
        assert "Failed to process the input audio data. Either" in response["message"]


# https://github.com/huggingface/transformers/commit/9c500015c556f9ddf6e7a7449d3f46b2e3ff8ea5
# caused a regression in beam search.
# https://github.com/huggingface/transformers/commit/a6b51e7341d702127a4a45f37439640840b5abf0
# fixed the regression but has not been released yet as of May 30, 2025.
@pytest.mark.skipif(
    Version("4.52.0") <= Version(transformers.__version__) < Version("4.53.0"),
    reason="Transformers 4.52 has a bug for beam search in whiper implementation",
)
def test_whisper_model_support_timestamps(whisper_pipeline):
    # Request payload to the model serving endpoint contains base64 encoded audio data
    audio = read_audio_data("bytes")
    encoded_audio = base64.b64encode(audio).decode("ascii")

    model_config = {
        "return_timestamps": "word",
        "chunk_length_s": 20,
        "stride_length_s": [5, 3],
    }

    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            whisper_pipeline,
            name="whisper_timestamps",
            model_config=model_config,
            input_example=(audio, model_config),
        )

    # Native transformers prediction as ground truth
    gt = whisper_pipeline(audio, **model_config)

    def _assert_prediction(pred):
        assert pred["text"] == gt["text"]
        assert len(pred["chunks"]) == len(gt["chunks"])
        for pred_chunk, gt_chunk in zip(pred["chunks"], gt["chunks"]):
            assert pred_chunk["text"] == gt_chunk["text"]
            # Timestamps are tuples, but converted to list when serialized to JSON.
            assert tuple(pred_chunk["timestamp"]) == gt_chunk["timestamp"]

    # Prediction with Pyfunc
    loaded_pyfunc = mlflow.pyfunc.load_model(model_info.model_uri)
    prediction = json.loads(loaded_pyfunc.predict(audio)[0])
    _assert_prediction(prediction)

    # Serve and score
    with pyfunc_scoring_endpoint(
        model_info.model_uri,
        extra_args=["--env-manager", "local"],
    ) as endpoint:
        content_type = pyfunc_scoring_server.CONTENT_TYPE_JSON
        payload = json.dumps({"inputs": [encoded_audio]})
        response = endpoint.invoke(payload, content_type=content_type)

        predictions = json.loads(response.content.decode("utf-8"))["predictions"]
        # When return_timestamps is specified, the predictions list contains json
        # serialized output from the pipeline.
        _assert_prediction(json.loads(predictions[0]))

        # Request with inference params
        payload = json.dumps({
            "inputs": [encoded_audio],
            "model_config": model_config,
        })
        response = endpoint.invoke(payload, content_type=content_type)
        predictions = json.loads(response.content.decode("utf-8"))["predictions"]
        _assert_prediction(json.loads(predictions[0]))


def test_whisper_model_pyfunc_with_malformed_input(whisper_pipeline, model_path):
    mlflow.transformers.save_model(
        transformers_model=whisper_pipeline,
        path=model_path,
        save_pretrained=False,
    )

    pyfunc_model = mlflow.pyfunc.load_model(model_path)

    invalid_audio = b"This isn't a real audio file"
    with pytest.raises(MlflowException, match="Failed to process the input audio data. Either"):
        pyfunc_model.predict([invalid_audio])

    bad_uri_msg = "An invalid string input was provided. String"

    with pytest.raises(MlflowException, match=bad_uri_msg):
        pyfunc_model.predict("An invalid path")

    with pytest.raises(MlflowException, match=bad_uri_msg):
        pyfunc_model.predict("//www.invalid.net/audio.wav")

    with pytest.raises(MlflowException, match=bad_uri_msg):
        pyfunc_model.predict("https:///my/audio.mp3")


@pytest.mark.parametrize("with_input_example", [True, False])
def test_audio_classification_pipeline(audio_classification_pipeline, with_input_example):
    audio = read_audio_data("bytes")

    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            audio_classification_pipeline,
            name="audio_classification",
            input_example=audio if with_input_example else None,
            save_pretrained=False,
        )

    inference_payload = json.dumps({"inputs": [base64.b64encode(audio).decode("ascii")]})

    response = pyfunc_serve_and_score_model(
        model_info.model_uri,
        data=inference_payload,
        content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON,
        extra_args=["--env-manager", "local"],
    )

    values = PredictionsResponse.from_json(response.content.decode("utf-8")).get_predictions()
    assert isinstance(values, pd.DataFrame)
    assert len(values) > 1
    assert list(values.columns) == ["score", "label"]


@pytest.mark.parametrize(
    "model_name",
    [
        "tiiuae/falcon-7b",
        "openai-community/gpt2",
        "PrunaAI/runwayml-stable-diffusion-v1-5-turbo-tiny-green-smashed",
    ],
)
def test_save_model_card_with_non_utf_characters(tmp_path, model_name):
    # non-ascii unicode characters
    test_text = (
        "Emoji testing! \u2728 \U0001f600 \U0001f609 \U0001f606 "
        "\U0001f970 \U0001f60e \U0001f917 \U0001f9d0"
    )

    card_data: ModelCard = huggingface_hub.ModelCard.load(model_name)
    card_data.text = card_data.text + "\n\n" + test_text
    custom_data = card_data.data.to_dict()
    custom_data["emojis"] = test_text

    card_data.data = huggingface_hub.CardData(**custom_data)
    _write_card_data(card_data, tmp_path)

    txt = tmp_path.joinpath(_CARD_TEXT_FILE_NAME).read_text()
    assert txt == card_data.text
    data = yaml.safe_load(tmp_path.joinpath(_CARD_DATA_FILE_NAME).read_text())
    assert data == card_data.data.to_dict()


def test_vision_pipeline_pyfunc_predict_with_kwargs(small_vision_model):
    artifact_path = "image_classification_model"

    parameters = {
        "top_k": 2,
    }
    inference_payload = json.dumps({
        "inputs": [image_url],
        "params": parameters,
    })

    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            small_vision_model,
            name=artifact_path,
            signature=infer_signature(
                image_url,
                mlflow.transformers.generate_signature_output(small_vision_model, image_url),
                params=parameters,
            ),
        )
        model_uri = model_info.model_uri
    transformers_loaded_model = mlflow.transformers.load_model(model_uri)
    expected_predictions = transformers_loaded_model.predict(image_url)

    response = pyfunc_serve_and_score_model(
        model_uri,
        data=inference_payload,
        content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON,
        extra_args=["--env-manager", "local"],
    )

    predictions = PredictionsResponse.from_json(response.content.decode("utf-8")).get_predictions()

    assert (
        list(predictions.to_dict("records")[0].values())
        == expected_predictions[: parameters["top_k"]]
    )


def test_qa_pipeline_pyfunc_predict_with_kwargs(small_qa_pipeline):
    artifact_path = "qa_model"
    data = {
        "question": [
            "What color is it?",
            "What does the 'wolf' howl at?",
        ],
        "context": [
            "Some people said it was green but I know that it's pink.",
            "The pack of 'wolves' stood on the cliff and a 'lone wolf' howled at "
            "the moon for hours.",
        ],
    }
    parameters = {
        "top_k": 2,
        "max_answer_len": 5,
    }
    inference_payload = json.dumps({
        "inputs": data,
        "params": parameters,
    })
    output = mlflow.transformers.generate_signature_output(small_qa_pipeline, data)
    signature_with_params = infer_signature(
        data,
        output,
        parameters,
    )
    expected_signature = ModelSignature(
        Schema([
            ColSpec(Array(DataType.string), name="question"),
            ColSpec(Array(DataType.string), name="context"),
        ]),
        Schema([ColSpec(DataType.string)]),
        ParamSchema([
            ParamSpec("top_k", DataType.long, 2),
            ParamSpec("max_answer_len", DataType.long, 5),
        ]),
    )
    assert signature_with_params == expected_signature

    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            small_qa_pipeline,
            name=artifact_path,
            signature=signature_with_params,
        )

    response = pyfunc_serve_and_score_model(
        model_info.model_uri,
        data=inference_payload,
        content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON,
        extra_args=["--env-manager", "local"],
    )
    values = PredictionsResponse.from_json(response.content.decode("utf-8")).get_predictions()

    assert values.to_dict(orient="records") == [
        {0: "pink"},
        {0: "pink."},
        {0: "the moon"},
        {0: "moon"},
    ]


def test_uri_directory_renaming_handling_pipeline(model_path, text_classification_pipeline):
    with mlflow.start_run():
        mlflow.transformers.save_model(
            transformers_model=text_classification_pipeline, path=model_path
        )

    absolute_model_directory = os.path.join(model_path, "model")
    renamed_to_old_convention = os.path.join(model_path, "pipeline")
    os.rename(absolute_model_directory, renamed_to_old_convention)

    # remove the 'model_binary' entries to emulate older versions of MLflow
    mlmodel_file = os.path.join(model_path, "MLmodel")
    with open(mlmodel_file) as yaml_file:
        mlmodel = yaml.safe_load(yaml_file)

    mlmodel["flavors"]["python_function"].pop("model_binary", None)
    mlmodel["flavors"]["transformers"].pop("model_binary", None)

    with open(mlmodel_file, "w") as yaml_file:
        yaml.safe_dump(mlmodel, yaml_file)

    loaded_model = mlflow.pyfunc.load_model(model_path)

    prediction = loaded_model.predict("test")
    assert isinstance(prediction, pd.DataFrame)
    assert isinstance(prediction["label"][0], str)


def test_uri_directory_renaming_handling_components(model_path, text_classification_pipeline):
    components = {
        "tokenizer": text_classification_pipeline.tokenizer,
        "model": text_classification_pipeline.model,
    }

    with mlflow.start_run():
        mlflow.transformers.save_model(transformers_model=components, path=model_path)

    absolute_model_directory = os.path.join(model_path, "model")
    renamed_to_old_convention = os.path.join(model_path, "pipeline")
    os.rename(absolute_model_directory, renamed_to_old_convention)

    # remove the 'model_binary' entries to emulate older versions of MLflow
    mlmodel_file = os.path.join(model_path, "MLmodel")
    with open(mlmodel_file) as yaml_file:
        mlmodel = yaml.safe_load(yaml_file)

    mlmodel["flavors"]["python_function"].pop("model_binary", None)
    mlmodel["flavors"]["transformers"].pop("model_binary", None)

    with open(mlmodel_file, "w") as yaml_file:
        yaml.safe_dump(mlmodel, yaml_file)

    loaded_model = mlflow.pyfunc.load_model(model_path)

    prediction = loaded_model.predict("test")
    assert isinstance(prediction, pd.DataFrame)
    assert isinstance(prediction["label"][0], str)


@skip_transformers_v5_or_later
def test_pyfunc_model_log_load_with_artifacts_snapshot():
    architecture = "prajjwal1/bert-tiny"
    tokenizer = transformers.AutoTokenizer.from_pretrained(architecture)
    model = transformers.BertForQuestionAnswering.from_pretrained(architecture)
    bert_tiny_pipeline = transformers.pipeline(
        task="question-answering", model=model, tokenizer=tokenizer
    )

    class QAModel(mlflow.pyfunc.PythonModel):
        def load_context(self, context):
            """
            This method initializes the tokenizer and language model
            using the specified snapshot location.
            """
            snapshot_location = context.artifacts["bert-tiny-model"].removeprefix("hf:/")
            # Initialize tokenizer and language model
            tokenizer = transformers.AutoTokenizer.from_pretrained(snapshot_location)
            model = transformers.BertForQuestionAnswering.from_pretrained(snapshot_location)
            self.pipeline = transformers.pipeline(
                task="question-answering", model=model, tokenizer=tokenizer
            )

        def predict(self, context, model_input, params=None):
            question = model_input["question"][0]
            if isinstance(question, np.ndarray):
                question = question.item()
            ctx = model_input["context"][0]
            if isinstance(ctx, np.ndarray):
                ctx = ctx.item()
            return self.pipeline(question=question, context=ctx)

    data = {"question": "Who's house?", "context": "The house is owned by Run."}
    pyfunc_artifact_path = "question_answering_model"
    with mlflow.start_run():
        model_info = mlflow.pyfunc.log_model(
            name=pyfunc_artifact_path,
            python_model=QAModel(),
            artifacts={"bert-tiny-model": "hf:/prajjwal1/bert-tiny"},
            input_example=data,
            signature=infer_signature(
                data, mlflow.transformers.generate_signature_output(bert_tiny_pipeline, data)
            ),
            extra_pip_requirements=["transformers", "torch", "numpy"],
        )

        pyfunc_model_path = _download_artifact_from_uri(model_info.model_uri)
        assert len(os.listdir(os.path.join(pyfunc_model_path, "artifacts"))) != 0
        model_config = Model.load(os.path.join(pyfunc_model_path, "MLmodel"))

    loaded_pyfunc_model = mlflow.pyfunc.load_model(model_uri=model_info.model_uri)
    assert model_config.to_yaml() == loaded_pyfunc_model.metadata.to_yaml()
    assert loaded_pyfunc_model.predict(data)["answer"] != ""

    # Test model serving
    inference_payload = json.dumps({"inputs": data})
    response = pyfunc_serve_and_score_model(
        model_info.model_uri,
        data=inference_payload,
        content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON,
        extra_args=["--env-manager", "local"],
    )
    values = PredictionsResponse.from_json(response.content.decode("utf-8")).get_predictions()

    assert values.to_dict(orient="records")[0]["answer"] != ""


def test_pyfunc_model_log_load_with_artifacts_snapshot_errors():
    class TestModel(mlflow.pyfunc.PythonModel):
        def predict(self, context, model_input, params=None):
            return model_input

    with mlflow.start_run():
        with pytest.raises(
            MlflowException,
            match=r"Failed to download snapshot from Hugging Face Hub "
            r"with artifact_uri: hf:/invalid-repo-id.",
        ):
            mlflow.pyfunc.log_model(
                name="pyfunc_artifact_path",
                python_model=TestModel(),
                artifacts={"some-model": "hf:/invalid-repo-id"},
            )


def test_model_distributed_across_devices():
    mock_model = mock.Mock()
    mock_model.device.type = "meta"
    mock_model.hf_device_map = {
        "layer1": mock.Mock(type="cpu"),
        "layer2": mock.Mock(type="cpu"),
        "layer3": mock.Mock(type="gpu"),
        "layer4": mock.Mock(type="disk"),
    }

    assert _is_model_distributed_in_memory(mock_model)


def test_model_on_single_device():
    mock_model = mock.Mock()
    mock_model.device.type = "cpu"
    mock_model.hf_device_map = {}

    assert not _is_model_distributed_in_memory(mock_model)


@skip_transformers_v5_or_later
def test_basic_model_with_accelerate_device_mapping_fails_save(tmp_path, model_path):
    task = "translation_en_to_de"
    architecture = "t5-small"
    model = transformers.T5ForConditionalGeneration.from_pretrained(
        pretrained_model_name_or_path=architecture,
        device_map={"shared": "cpu", "encoder": "cpu", "decoder": "disk", "lm_head": "disk"},
        offload_folder=str(tmp_path / "weights"),
        low_cpu_mem_usage=True,
    )

    tokenizer = transformers.T5TokenizerFast.from_pretrained(
        pretrained_model_name_or_path=architecture, model_max_length=100
    )
    pipeline = transformers.pipeline(task=task, model=model, tokenizer=tokenizer)

    with pytest.raises(
        MlflowException,
        match="The model that is attempting to be saved has been loaded into memory",
    ):
        mlflow.transformers.save_model(transformers_model=pipeline, path=model_path)


@pytest.mark.skipif(
    Version(transformers.__version__) > Version("4.44.2"),
    reason="Multi-task pipeline (t5) has a loading issue with Transformers 4.45.x. "
    "See https://github.com/huggingface/transformers/issues/33398 for more details.",
)
def test_basic_model_with_accelerate_homogeneous_mapping_works(model_path):
    task = "translation_en_to_de"
    architecture = "t5-small"
    model = transformers.T5ForConditionalGeneration.from_pretrained(
        pretrained_model_name_or_path=architecture,
        device_map={"shared": "cpu", "encoder": "cpu", "decoder": "cpu", "lm_head": "cpu"},
        low_cpu_mem_usage=True,
    )

    tokenizer = transformers.T5TokenizerFast.from_pretrained(
        pretrained_model_name_or_path=architecture, model_max_length=100
    )
    pipeline = transformers.pipeline(task=task, model=model, tokenizer=tokenizer)

    mlflow.transformers.save_model(transformers_model=pipeline, path=model_path)

    loaded = mlflow.transformers.load_model(model_path)
    text = "Apples are delicious"
    assert loaded(text) == pipeline(text)


def test_qa_model_model_size_bytes(small_qa_pipeline, tmp_path):
    def _calculate_expected_size(path_or_dir):
        # this helper function does not consider subdirectories
        expected_size = 0
        if path_or_dir.is_dir():
            for path in path_or_dir.iterdir():
                if not path.is_file():
                    continue
                expected_size += path.stat().st_size
        elif path_or_dir.is_file():
            expected_size = path_or_dir.stat().st_size
        return expected_size

    mlflow.transformers.save_model(
        transformers_model=small_qa_pipeline,
        path=tmp_path,
    )

    # expected size only counts for files saved before the MLmodel file is saved
    model_dir = tmp_path.joinpath("model")
    tokenizer_dir = tmp_path.joinpath("components").joinpath("tokenizer")
    expected_size = 0
    for folder in [model_dir, tokenizer_dir]:
        expected_size += _calculate_expected_size(folder)
    other_files = ["model_card.md", "model_card_data.yaml", "LICENSE.txt"]
    for file in other_files:
        path = tmp_path.joinpath(file)
        expected_size += _calculate_expected_size(path)

    mlmodel = yaml.safe_load(tmp_path.joinpath("MLmodel").read_bytes())
    assert mlmodel["model_size_bytes"] == expected_size


@pytest.mark.parametrize(
    ("task", "input_example"),
    [
        ("llm/v1/completions", None),
        ("llm/v1/chat", None),
        (
            "llm/v1/completions",
            {
                "prompt": "How to learn Python in 3 weeks?",
                "max_tokens": 10,
                "stop": "Python",
            },
        ),
        (
            "llm/v1/chat",
            {
                "messages": [
                    {"role": "system", "content": "Hello, how are you?"},
                ],
                "temperature": 0.5,
                "max_tokens": 50,
            },
        ),
    ],
)
def test_text_generation_save_model_with_inference_task(
    monkeypatch, task, input_example, text_generation_pipeline, model_path
):
    # Strictly raise error during requirements inference for testing purposes
    monkeypatch.setenv("MLFLOW_REQUIREMENTS_INFERENCE_RAISE_ERRORS", "true")

    mlflow.transformers.save_model(
        transformers_model=text_generation_pipeline,
        path=model_path,
        task=task,
        input_example=input_example,
    )

    mlmodel = yaml.safe_load(model_path.joinpath("MLmodel").read_bytes())
    flavor_config = mlmodel["flavors"]["transformers"]
    assert flavor_config["inference_task"] == task
    assert mlmodel["metadata"]["task"] == task

    if input_example:
        saved_input_example = json.loads(model_path.joinpath("input_example.json").read_text())
        assert saved_input_example == input_example


def test_text_generation_save_model_with_invalid_inference_task(
    text_generation_pipeline, model_path
):
    with pytest.raises(
        MlflowException, match=r"The task provided is invalid.*Must be.*llm/v1/completions"
    ):
        mlflow.transformers.save_model(
            transformers_model=text_generation_pipeline,
            path=model_path,
            task="llm/v1/invalid",
        )


def test_text_generation_log_model_with_mismatched_task(text_generation_pipeline):
    with pytest.raises(
        MlflowException, match=r"LLM v1 task type 'llm/v1/chat' is specified in metadata, but"
    ):
        with mlflow.start_run():
            mlflow.transformers.log_model(
                text_generation_pipeline,
                name="model",
                # Task argument and metadata task are different
                task=None,
                metadata={"task": "llm/v1/chat"},
            )


def test_text_generation_task_completions_predict_with_max_tokens(
    text_generation_pipeline, model_path
):
    mlflow.transformers.save_model(
        transformers_model=text_generation_pipeline,
        path=model_path,
        task="llm/v1/completions",
        model_config={"max_tokens": 10},
    )

    pyfunc_loaded = mlflow.pyfunc.load_model(model_path)

    inference = pyfunc_loaded.predict(
        {"prompt": "How to learn Python in 3 weeks?", "max_tokens": 10},
    )

    assert isinstance(inference[0], dict)
    assert inference[0]["model"] == "distilgpt2"
    assert inference[0]["object"] == "text_completion"
    assert (
        inference[0]["choices"][0]["finish_reason"] == "length"
        and inference[0]["usage"]["completion_tokens"] == 10
    ) or (
        inference[0]["choices"][0]["finish_reason"] == "stop"
        and inference[0]["usage"]["completion_tokens"] < 10
    )

    # Override model_config with runtime params
    inference = pyfunc_loaded.predict(
        {"prompt": "How to learn Python in 3 weeks?", "max_tokens": 5},
    )
    assert 6 > inference[0]["usage"]["completion_tokens"] > 0


def test_text_generation_task_completions_predict_with_stop(text_generation_pipeline, model_path):
    mlflow.transformers.save_model(
        transformers_model=text_generation_pipeline,
        path=model_path,
        task="llm/v1/completions",
        metadata={"foo": "bar"},
        model_config={"stop": ["Python"], "max_tokens": 50},
    )

    pyfunc_loaded = mlflow.pyfunc.load_model(model_path)
    inference = pyfunc_loaded.predict(
        {"prompt": "How to learn Python in 3 weeks?"},
    )

    if "Python" not in inference[0]["choices"][0]["text"]:
        pytest.skip(
            "Model did not generate text containing 'Python', "
            "skipping validation of stop parameter in inference"
        )

    assert (
        inference[0]["choices"][0]["finish_reason"] == "stop"
        and inference[0]["usage"]["completion_tokens"] < 50
    ) or (
        inference[0]["choices"][0]["finish_reason"] == "length"
        and inference[0]["usage"]["completion_tokens"] == 50
    )

    assert inference[0]["choices"][0]["text"].endswith("Python")

    # Override model_config with runtime params
    inference = pyfunc_loaded.predict(
        {"prompt": "How to learn Python in 3 weeks?", "stop": ["Abracadabra"]},
    )

    response_text = inference[0]["choices"][0]["text"]

    # Only check for early stopping if we stop on the word "Python".
    # If we exhaust the token limit, there is a non-insignificant chance of
    # terminating on the word due to max tokens, which should not count as
    # a stop word abort if there are multiple instances of the word in the text.
    if 0 < response_text.count("Python") < 2:
        assert not inference[0]["choices"][0]["text"].endswith("Python")


def test_text_generation_task_completions_serve(text_generation_pipeline):
    data = {"prompt": "How to learn Python in 3 weeks?"}

    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            text_generation_pipeline,
            name="model",
            task="llm/v1/completions",
        )

    inference_payload = json.dumps({"inputs": data})

    response = pyfunc_serve_and_score_model(
        model_info.model_uri,
        data=inference_payload,
        content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON,
        extra_args=["--env-manager", "local"],
    )
    values = PredictionsResponse.from_json(response.content.decode("utf-8")).get_predictions()
    output_dict = values.to_dict("records")[0]
    assert output_dict["choices"][0]["text"] is not None
    assert output_dict["choices"][0]["finish_reason"] == "stop"
    assert output_dict["usage"]["prompt_tokens"] < 20


def test_llm_v1_task_embeddings_predict(feature_extraction_pipeline, model_path):
    mlflow.transformers.save_model(
        transformers_model=feature_extraction_pipeline,
        path=model_path,
        input_examples=["Football", "Soccer"],
        task="llm/v1/embeddings",
    )

    mlmodel = yaml.safe_load(model_path.joinpath("MLmodel").read_bytes())

    flavor_config = mlmodel["flavors"]["transformers"]
    assert flavor_config["inference_task"] == "llm/v1/embeddings"
    assert mlmodel["metadata"]["task"] == "llm/v1/embeddings"

    pyfunc_loaded = mlflow.pyfunc.load_model(model_path)

    # Predict with single string input
    prediction = pyfunc_loaded.predict({"input": "A great day"})
    assert prediction["object"] == "list"
    assert len(prediction["data"]) == 1
    assert prediction["data"][0]["object"] == "embedding"
    assert prediction["usage"]["prompt_tokens"] == 5
    assert len(prediction["data"][0]["embedding"]) == 384

    # Predict with list of string input
    prediction = pyfunc_loaded.predict({"input": ["A great day", "A bad day"]})
    assert prediction["object"] == "list"
    assert len(prediction["data"]) == 2
    assert prediction["data"][0]["object"] == "embedding"
    assert prediction["usage"]["prompt_tokens"] == 10
    assert len(prediction["data"][0]["embedding"]) == 384

    # Predict with pandas dataframe input
    df = pd.DataFrame({"input": ["A great day", "A bad day", "A good day"]})
    prediction = pyfunc_loaded.predict(df)
    assert prediction["object"] == "list"
    assert len(prediction["data"]) == 3
    assert prediction["data"][0]["object"] == "embedding"
    assert prediction["usage"]["prompt_tokens"] == 15
    assert len(prediction["data"][0]["embedding"]) == 384


@pytest.mark.parametrize(
    "request_payload",
    [
        {"input": "A single string"},
        {
            "inputs": {"input": ["A list of strings"]},
        },
    ],
)
def test_llm_v1_task_embeddings_serve(feature_extraction_pipeline, request_payload):
    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            feature_extraction_pipeline,
            name="model",
            input_examples=["Football", "Soccer"],
            task="llm/v1/embeddings",
        )

    response = pyfunc_serve_and_score_model(
        model_info.model_uri,
        data=json.dumps(request_payload),
        content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON,
        extra_args=["--env-manager", "local"],
    )
    response = json.loads(response.content.decode("utf-8"))
    prediction = response["predictions"] if "inputs" in request_payload else response

    assert prediction["object"] == "list"
    assert len(prediction["data"]) == 1
    assert prediction["data"][0]["object"] == "embedding"
    assert len(prediction["data"][0]["embedding"]) == 384


def test_get_task_for_model():
    with mock.patch("transformers.pipelines.get_task") as mock_get_task:
        mock_get_task.return_value = "feature-extraction"
        assert _get_task_for_model("model") == "feature-extraction"

        # Some model task is not supported by Transformers pipeline yet. Then fall back
        # to the default task if provided, otherwise raise an exception.
        mock_get_task.return_value = "unsupported-task"
        assert (
            _get_task_for_model("model", default_task="feature-extraction") == "feature-extraction"
        )

        with pytest.raises(MlflowException, match="Cannot construct transformers pipeline"):
            _get_task_for_model("model")

        # If get_task raises an exception, fall back to the default task if provided.
        mock_get_task.side_effect = RuntimeError("Some error")
        assert (
            _get_task_for_model("model", default_task="feature-extraction") == "feature-extraction"
        )

        with pytest.raises(MlflowException, match="The task could not be inferred"):
            _get_task_for_model("model")


@skip_transformers_v5_or_later
def test_local_custom_model_save_and_load(text_generation_pipeline, model_path, tmp_path):
    local_repo_path = tmp_path / "local_repo"
    text_generation_pipeline.save_pretrained(local_repo_path)

    locally_loaded_model = transformers.AutoModelForCausalLM.from_pretrained(local_repo_path)
    tokenizer = transformers.AutoTokenizer.from_pretrained(
        local_repo_path, chat_template=CHAT_TEMPLATE
    )
    model_dict = {"model": locally_loaded_model, "tokenizer": tokenizer}

    # 1. Save local custom model without specifying task -> raises MlflowException
    with pytest.raises(MlflowException, match=r"The task could not be inferred"):
        mlflow.transformers.save_model(transformers_model=model_dict, path=model_path)

    # 2. Save local custom model with task -> saves successfully
    mlflow.transformers.save_model(
        transformers_model=model_dict,
        path=model_path,
        task="text-generation",
    )

    pyfunc_loaded = mlflow.pyfunc.load_model(model_path)

    inference = pyfunc_loaded.predict("How to save Transformer model?")
    assert isinstance(inference[0], str)
    assert inference[0].startswith("How to save Transformer model?")

    # 3. Save local custom model with LLM v1 chat inference task -> saves successfully
    #    with the corresponding Transformers task
    shutil.rmtree(model_path)

    mlflow.transformers.save_model(
        transformers_model=model_dict,
        path=model_path,
        task="llm/v1/chat",
    )

    mlmodel = yaml.safe_load(model_path.joinpath("MLmodel").read_bytes())
    flavor_config = mlmodel["flavors"]["transformers"]
    assert flavor_config["task"] == "text-generation"
    assert flavor_config["inference_task"] == "llm/v1/chat"
    assert mlmodel["metadata"]["task"] == "llm/v1/chat"

    pyfunc_loaded = mlflow.pyfunc.load_model(model_path)

    inference = pyfunc_loaded.predict({
        "messages": [
            {
                "role": "user",
                "content": "How to save Transformer model?",
            }
        ]
    })
    assert isinstance(inference[0], dict)
    assert inference[0]["choices"][0]["message"]["role"] == "assistant"


def test_model_config_is_not_mutated_after_prediction(text2text_generation_pipeline):
    model_config = {
        "top_k": 2,
        "num_beams": 5,
        "max_length": 30,
        "max_new_tokens": 500,
    }

    # Params will be used to override the values of model_config but should not mutate it
    params = {
        "top_k": 30,
        "max_length": 500,
        "max_new_tokens": 5,
    }

    pyfunc_model = _TransformersWrapper(text2text_generation_pipeline, model_config=model_config)
    assert pyfunc_model.model_config["top_k"] == 2

    prediction_output = pyfunc_model.predict(
        "rocket moon ship astronaut space gravity", params=params
    )

    assert pyfunc_model.model_config["top_k"] == 2
    assert pyfunc_model.model_config["num_beams"] == 5
    assert pyfunc_model.model_config["max_length"] == 30
    assert pyfunc_model.model_config["max_new_tokens"] == 500
    assert len(prediction_output[0].split(" ")) <= 5


def test_text_generation_task_chat_predict(text_generation_pipeline, model_path):
    mlflow.transformers.save_model(
        transformers_model=text_generation_pipeline,
        path=model_path,
        task="llm/v1/chat",
    )

    pyfunc_loaded = mlflow.pyfunc.load_model(model_path)

    inference = pyfunc_loaded.predict({
        "messages": [
            {"role": "system", "content": "Hello, how can I help you today?"},
            {"role": "user", "content": "How to learn Python in 3 weeks?"},
        ],
        "max_tokens": 10,
    })

    assert inference[0]["choices"][0]["message"]["role"] == "assistant"
    assert (
        inference[0]["choices"][0]["finish_reason"] == "length"
        and inference[0]["usage"]["completion_tokens"] == 10
    ) or (
        inference[0]["choices"][0]["finish_reason"] == "stop"
        and inference[0]["usage"]["completion_tokens"] < 10
    )


def test_text_generation_task_chat_serve(text_generation_pipeline):
    data = {
        "messages": [
            {"role": "user", "content": "How to learn Python in 3 weeks?"},
        ],
        "max_tokens": 10,
    }

    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            text_generation_pipeline,
            name="model",
            task="llm/v1/chat",
        )

    inference_payload = json.dumps(data)

    response = pyfunc_serve_and_score_model(
        model_info.model_uri,
        data=inference_payload,
        content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON,
        extra_args=["--env-manager", "local"],
    )

    output_dict = json.loads(response.content)[0]
    assert output_dict["choices"][0]["message"] is not None
    assert (
        output_dict["choices"][0]["finish_reason"] == "length"
        and output_dict["usage"]["completion_tokens"] == 10
    ) or (
        output_dict["choices"][0]["finish_reason"] == "stop"
        and output_dict["usage"]["completion_tokens"] < 10
    )
    assert output_dict["usage"]["prompt_tokens"] < 20


HF_COMMIT_HASH_PATTERN = re.compile(r"^[a-z0-9]{40}$")


@pytest.mark.parametrize(
    ("model_fixture", "input_example", "components"),
    [
        ("text2text_generation_pipeline", "What is MLflow?", {"tokenizer"}),
        ("text_generation_pipeline", "What is MLflow?", {"tokenizer"}),
        (
            "small_vision_model",
            image_url,
            {"image_processor"} if IS_NEW_FEATURE_EXTRACTION_API else {"feature_extractor"},
        ),
        (
            "component_multi_modal",
            {"text": "What is MLflow?", "image": image_url},
            {"image_processor", "tokenizer"}
            if IS_NEW_FEATURE_EXTRACTION_API
            else {"feature_extractor", "tokenizer"},
        ),
        ("fill_mask_pipeline", "The quick brown <mask> jumps over the lazy dog.", {"tokenizer"}),
        ("whisper_pipeline", lambda: read_audio_data("bytes"), {"feature_extractor", "tokenizer"}),
        ("feature_extraction_pipeline", "What is MLflow?", {"tokenizer"}),
    ],
)
def test_save_and_load_pipeline_without_save_pretrained_false(
    model_fixture, input_example, components, model_path, request
):
    pipeline = request.getfixturevalue(model_fixture)
    model = pipeline["model"] if isinstance(pipeline, dict) else pipeline.model

    mlflow.transformers.save_model(
        transformers_model=pipeline,
        path=model_path,
        save_pretrained=False,
    )

    # No weights should be saved
    assert not model_path.joinpath("model").exists()
    assert not model_path.joinpath("components").exists()

    # Validate the contents of MLModel file
    mlmodel = Model.load(str(model_path.joinpath("MLmodel")))
    flavor_conf = mlmodel.flavors["transformers"]
    assert "model_binary" not in flavor_conf
    assert flavor_conf["source_model_name"] == model.name_or_path
    assert HF_COMMIT_HASH_PATTERN.match(flavor_conf["source_model_revision"])
    assert set(flavor_conf["components"]) == components
    for c in components:
        component = pipeline[c] if isinstance(pipeline, dict) else getattr(pipeline, c)
        assert flavor_conf[f"{c}_name"] == getattr(component, "name_or_path", model.name_or_path)
        assert HF_COMMIT_HASH_PATTERN.match(flavor_conf[f"{c}_revision"])

    # Validate pyfunc load and prediction (if pyfunc supported)
    if "python_function" in mlmodel.flavors:
        if callable(input_example):
            input_example = input_example()
        mlflow.pyfunc.load_model(model_path).predict(input_example)


# Patch tempdir just to verify the invocation
def test_persist_pretrained_model(small_qa_pipeline):
    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            small_qa_pipeline,
            name="model",
            save_pretrained=False,
            pip_requirements=["mlflow"],  # For speed up logging
        )

    artifact_path = Path(mlflow.artifacts.download_artifacts(model_info.model_uri))
    model_path = artifact_path / "model"
    tokenizer_path = artifact_path / "components" / "tokenizer"

    original_config = Model.load(artifact_path).flavors["transformers"]
    assert "model_binary" not in original_config
    assert "source_model_revision" in original_config
    assert not model_path.exists()
    assert not tokenizer_path.exists()

    with mock.patch(
        "mlflow.transformers.TempDir", side_effect=mlflow.utils.file_utils.TempDir
    ) as mock_tmpdir:
        mlflow.transformers.persist_pretrained_model(model_info.model_uri)
        mock_tmpdir.assert_called_once()

    updated_config = Model.load(model_info.model_uri).flavors["transformers"]
    assert "model_binary" in updated_config
    assert "source_model_revision" not in updated_config
    assert model_path.exists()
    model_path_files = list(model_path.iterdir())
    assert len(model_path_files) > 0
    assert tokenizer_path.exists()
    assert (tokenizer_path / "tokenizer.json").exists()

    # Repeat persisting the model will no-op
    with mock.patch(
        "mlflow.transformers.TempDir", side_effect=mlflow.utils.file_utils.TempDir
    ) as mock_tmpdir:
        mlflow.transformers.persist_pretrained_model(model_info.model_uri)
        mock_tmpdir.assert_not_called()


def test_small_qa_pipeline_copy_metadata_in_databricks(
    mock_is_in_databricks, small_qa_pipeline, tmp_path
):
    artifact_path = "transformers"
    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            small_qa_pipeline,
            name=artifact_path,
        )
    artifact_path = mlflow.artifacts.download_artifacts(
        artifact_uri=model_info.model_uri, dst_path=tmp_path.as_posix()
    )

    # Metadata should be copied only in Databricks
    metadata_path = os.path.join(artifact_path, "metadata")
    if mock_is_in_databricks.return_value:
        assert set(os.listdir(metadata_path)) == set(METADATA_FILES)
    else:
        assert not os.path.exists(metadata_path)
    mock_is_in_databricks.assert_called_once()


def test_peft_pipeline_copy_metadata_in_databricks(mock_is_in_databricks, peft_pipeline, tmp_path):
    artifact_path = "transformers"
    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            peft_pipeline,
            name=artifact_path,
        )

    artifact_path = mlflow.artifacts.download_artifacts(
        artifact_uri=model_info.model_uri, dst_path=tmp_path.as_posix()
    )

    # Metadata should be copied only in Databricks
    metadata_path = os.path.join(artifact_path, "metadata")
    if mock_is_in_databricks.return_value:
        assert set(os.listdir(metadata_path)) == set(METADATA_FILES)
    else:
        assert not os.path.exists(metadata_path)
    mock_is_in_databricks.assert_called_once()


@pytest.mark.parametrize("device", ["cpu", "cuda", 0, -1, None])
def test_device_param_on_load_model(device, small_qa_pipeline, model_path, monkeypatch):
    mlflow.transformers.save_model(small_qa_pipeline, path=model_path)
    conf = mlflow.transformers.load_model(model_path, return_type="components", device=device)
    assert conf.get("device") == device

    monkeypatch.setenv("MLFLOW_HUGGINGFACE_USE_DEVICE_MAP", "true")
    if device is None:
        conf = mlflow.transformers.load_model(model_path, return_type="components", device=device)
        assert conf.get("device") is None
    else:
        with pytest.raises(
            MlflowException,
            match=r"The environment variable MLFLOW_HUGGINGFACE_USE_DEVICE_MAP is set to True, "
            rf"but the `device` argument is provided with value {device}.",
        ):
            mlflow.transformers.load_model(model_path, return_type="components", device=device)


@pytest.fixture
def local_checkpoint_path(tmp_path):
    """
    Fixture to create a local model checkpoint for testing fine-tuning scenario.
    """
    model = transformers.AutoModelForCausalLM.from_pretrained("distilgpt2")

    class DummyDataset(torch.utils.data.Dataset):
        def __getitem__(self, idx):
            pass

        def __len__(self):
            return 1

    # Create a trainer and save model, but not running the actual training
    training_args = transformers.TrainingArguments(
        output_dir=tmp_path / "result",
        num_train_epochs=1,
        per_device_train_batch_size=4,
        report_to="none",
    )
    trainer = transformers.Trainer(model=model, args=training_args, train_dataset=DummyDataset())

    checkpoint_path = tmp_path / "checkpoint"
    trainer.save_model(checkpoint_path)

    # The tokenizer should also be saved in the checkpoint
    tokenizer = transformers.AutoTokenizer.from_pretrained(
        # Chat template is required to test with llm/v1/chat task
        "distilgpt2",
        chat_template=CHAT_TEMPLATE,
    )
    tokenizer.save_pretrained(checkpoint_path)

    return str(checkpoint_path)


def test_save_model_from_local_checkpoint(model_path, local_checkpoint_path):
    with mock.patch("mlflow.transformers._logger") as mock_logger:
        mlflow.transformers.save_model(
            transformers_model=local_checkpoint_path,
            task="text-generation",
            path=model_path,
            input_example=["What is MLflow?"],
        )

    logged_info = Model.load(model_path)
    flavor_conf = logged_info.flavors["transformers"]
    assert flavor_conf["source_model_name"] == local_checkpoint_path
    assert flavor_conf["task"] == "text-generation"
    if not IS_TRANSFORMERS_V5_OR_LATER:
        assert flavor_conf["framework"] == "pt"
    assert flavor_conf["instance_type"] == "TextGenerationPipeline"
    expected_tokenizer_type = (
        "GPT2Tokenizer" if IS_TRANSFORMERS_V5_OR_LATER else "GPT2TokenizerFast"
    )
    assert flavor_conf["tokenizer_type"] == expected_tokenizer_type

    # Default task signature should be used
    assert logged_info.signature.inputs == Schema([ColSpec(DataType.string)])
    assert logged_info.signature.outputs == Schema([ColSpec(DataType.string)])

    # Default requirements should be used
    info_calls = mock_logger.info.call_args_list
    assert any("A local checkpoint path or PEFT model" in c[0][0] for c in info_calls)
    with model_path.joinpath("requirements.txt").open() as f:
        reqs = {req.split("==")[0] for req in f.read().split("\n")}
    assert reqs == {"mlflow", "accelerate", "transformers", "torch", "torchvision"}

    # Load as native pipeline
    loaded_pipeline = mlflow.transformers.load_model(model_path)
    assert isinstance(loaded_pipeline, transformers.TextGenerationPipeline)

    query = "What is MLflow?"
    pred_native = loaded_pipeline(query)[0]
    assert pred_native["generated_text"].startswith(query)

    # Load as pyfunc
    loaded_pyfunc = mlflow.pyfunc.load_model(model_path)
    pred_pyfunc = loaded_pyfunc.predict(query)[0]
    assert pred_pyfunc.startswith(query)

    # Serve
    response = pyfunc_serve_and_score_model(
        model_path,
        data=json.dumps({"inputs": [query]}),
        content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON,
        extra_args=["--env-manager", "local"],
    )
    pred_serve = json.loads(response.content.decode("utf-8"))
    assert pred_serve["predictions"][0].startswith(query)


@skip_transformers_v5_or_later
def test_save_model_from_local_checkpoint_with_custom_tokenizer(model_path, local_checkpoint_path):
    # When a custom tokenizer is also saved in the checkpoint, MLflow should save and load it.
    tokenizer = transformers.AutoTokenizer.from_pretrained("distilroberta-base")
    tokenizer.add_special_tokens({"additional_special_tokens": ["<sushi>"]})
    tokenizer.save_pretrained(local_checkpoint_path)

    mlflow.transformers.save_model(
        transformers_model=local_checkpoint_path,
        path=model_path,
        task="text-generation",
        input_example=["What is MLflow?"],
    )

    # The custom tokenizer should be loaded
    loaded_pipeline = mlflow.transformers.load_model(model_path)
    tokenizer = loaded_pipeline.tokenizer
    assert tokenizer.special_tokens_map["additional_special_tokens"] == ["<sushi>"]


def test_save_model_from_local_checkpoint_with_llm_inference_task(
    model_path, local_checkpoint_path
):
    mlflow.transformers.save_model(
        transformers_model=local_checkpoint_path,
        path=model_path,
        task="llm/v1/chat",
        input_example=["What is MLflow?"],
    )

    logged_info = Model.load(model_path)
    flavor_conf = logged_info.flavors["transformers"]
    assert flavor_conf["source_model_name"] == local_checkpoint_path
    assert flavor_conf["task"] == "text-generation"
    assert flavor_conf["inference_task"] == "llm/v1/chat"

    # Load as pyfunc
    loaded_pyfunc = mlflow.pyfunc.load_model(model_path)
    response = loaded_pyfunc.predict({
        "messages": [
            {"role": "system", "content": "Hello, how can I help you today?"},
            {"role": "user", "content": "What is MLflow?"},
        ],
    })
    assert response[0]["choices"][0]["message"]["role"] == "assistant"
    assert response[0]["choices"][0]["message"]["content"] is not None


def test_save_model_from_local_checkpoint_invalid_arguments(model_path, local_checkpoint_path):
    with pytest.raises(MlflowException, match=r"The `task` argument must be specified"):
        mlflow.transformers.save_model(
            transformers_model=local_checkpoint_path,
            path=model_path,
        )

    with pytest.raises(
        MlflowException, match=r"The `save_pretrained` argument must be set to True"
    ):
        mlflow.transformers.save_model(
            transformers_model=local_checkpoint_path,
            path=model_path,
            task="fill-mask",
            save_pretrained=False,
        )

    with pytest.raises(
        MlflowException,
        match=r"The provided directory invalid path does not contain a config.json file.",
    ):
        mlflow.transformers.save_model(
            transformers_model="invalid path",
            path=model_path,
            task="fill-mask",
        )


@pytest.mark.parametrize(
    ("model_fixture", "should_skip_validation"),
    [
        ("local_checkpoint_path", True),
        ("fill_mask_pipeline", False),
    ],
)
def test_log_model_skip_validating_serving_input_for_local_checkpoint(
    model_fixture,
    should_skip_validation,
    tmp_path,
    request,
):
    # input to avoid expensive computation
    model = request.getfixturevalue(model_fixture)
    with mock.patch("mlflow.models.utils._validate_serving_input") as mock_validate_input:
        # Ensure mlflow skips serving input validation for local checkpoint
        with mlflow.start_run():
            model_info = mlflow.transformers.log_model(
                model,
                name="model",
                task="fill-mask",
                input_example=["How are you?"],
            )

    # Serving input should exist regardless of the skip validation
    mlflow_model = Model.load(model_info.model_uri)
    local_path = _download_artifact_from_uri(model_info.model_uri, output_path=tmp_path)
    serving_input = mlflow_model.get_serving_input(local_path)
    assert json.loads(serving_input) == {"inputs": ["How are you?"]}

    if should_skip_validation:
        mock_validate_input.assert_not_called()
    else:
        mock_validate_input.assert_called_once()