Back to Repositories

Validating ML Pipeline Integration Workflows in gradio-app/gradio

This test suite validates the integration of Hugging Face Transformers and Diffusers pipelines with the Gradio interface. It ensures proper handling of various ML pipeline types including text generation, image processing, and question answering systems.

Test Coverage Overview

The test suite provides comprehensive coverage of pipeline integrations:

Transformers pipelines including text-to-text, audio classification, speech recognition, and various NLP tasks
Diffusers pipelines covering stable diffusion variants for image generation and manipulation
Edge cases for unsupported pipeline types and error handling
Integration testing between Gradio interface components and pipeline outputs

Implementation Analysis

The testing approach uses a combination of pytest and unittest frameworks with mock objects for pipeline simulation. The implementation follows a structured pattern testing both the pipeline handler functions and interface generation, with specific focus on input/output component validation and label verification.

Technical Details

Testing Frameworks: pytest, unittest
Mock Objects: unittest.mock.MagicMock
Key Libraries: transformers, diffusers, gradio
Test Decorators: @pytest.mark.flaky for handling non-deterministic behavior
Assert Patterns: Type checking, component validation, error handling

Best Practices Demonstrated

The test suite exemplifies several testing best practices:

Systematic organization of test cases by pipeline type
Proper isolation of test cases using mock objects
Comprehensive error case handling
Clear test method naming conventions
Efficient test setup and teardown patterns

gradio-app/gradio

test/test_pipelines.py

            
import unittest
from unittest.mock import MagicMock

import pytest
import transformers
from diffusers import (
    StableDiffusionDepth2ImgPipeline,  # type: ignore
    StableDiffusionImageVariationPipeline,  # type: ignore
    StableDiffusionImg2ImgPipeline,  # type: ignore
    StableDiffusionInpaintPipeline,  # type: ignore
    StableDiffusionInstructPix2PixPipeline,  # type: ignore
    StableDiffusionPipeline,  # type: ignore
    StableDiffusionUpscalePipeline,  # type: ignore
)
from transformers import (
    AudioClassificationPipeline,
    AutomaticSpeechRecognitionPipeline,
    DocumentQuestionAnsweringPipeline,
    FeatureExtractionPipeline,
    FillMaskPipeline,
    ImageClassificationPipeline,
    ImageToTextPipeline,
    ObjectDetectionPipeline,
    QuestionAnsweringPipeline,
    SummarizationPipeline,
    Text2TextGenerationPipeline,
    TextClassificationPipeline,
    TextGenerationPipeline,
    TranslationPipeline,
    VisualQuestionAnsweringPipeline,
    ZeroShotClassificationPipeline,
)

import gradio as gr
from gradio.pipelines_utils import (
    handle_diffusers_pipeline,
    handle_transformers_pipeline,
)


@pytest.mark.flaky
def test_text_to_text_model_from_pipeline():
    pipe = transformers.pipeline(model="sshleifer/bart-tiny-random")
    io = gr.Interface.from_pipeline(pipe)
    output = io("My name is Sylvain and I work at Hugging Face in Brooklyn")
    assert isinstance(output, str)


@pytest.mark.flaky
def test_stable_diffusion_pipeline():
    pipe = StableDiffusionPipeline.from_pretrained("hf-internal-testing/tiny-sd-pipe")
    io = gr.Interface.from_pipeline(pipe)
    output = io("An astronaut", "low quality", 3, 7.5)
    assert isinstance(output, str)


@pytest.mark.flaky
def test_interface_in_blocks():
    pipe1 = transformers.pipeline(model="sshleifer/bart-tiny-random")
    pipe2 = transformers.pipeline(model="sshleifer/bart-tiny-random")
    with gr.Blocks() as demo:
        with gr.Tab("Image Inference"):
            gr.Interface.from_pipeline(pipe1)
        with gr.Tab("Image Inference"):
            gr.Interface.from_pipeline(pipe2)
    demo.launch(prevent_thread_lock=True)
    demo.close()


@pytest.mark.flaky
def test_transformers_load_from_pipeline():
    from transformers import pipeline

    pipe = pipeline(model="deepset/roberta-base-squad2")
    io = gr.Interface.from_pipeline(pipe)
    assert io.input_components[0].label == "Context"
    assert io.input_components[1].label == "Question"
    assert io.output_components[0].label == "Answer"
    assert io.output_components[1].label == "Score"


class TestHandleTransformersPipelines(unittest.TestCase):
    def test_audio_classification_pipeline(self):
        pipe = MagicMock(spec=AudioClassificationPipeline)
        pipeline_info = handle_transformers_pipeline(pipe)
        assert pipeline_info is not None
        assert pipeline_info["inputs"].label == "Input"
        assert pipeline_info["outputs"].label == "Class"

    def test_automatic_speech_recognition_pipeline(self):
        pipe = MagicMock(spec=AutomaticSpeechRecognitionPipeline)
        pipeline_info = handle_transformers_pipeline(pipe)
        assert pipeline_info is not None
        assert pipeline_info["inputs"].label == "Input"
        assert pipeline_info["outputs"].label == "Output"

    def test_object_detection_pipeline(self):
        pipe = MagicMock(spec=ObjectDetectionPipeline)
        pipeline_info = handle_transformers_pipeline(pipe)
        assert pipeline_info is not None
        assert pipeline_info["inputs"].label == "Input Image"
        assert pipeline_info["outputs"].label == "Objects Detected"

    def test_feature_extraction_pipeline(self):
        pipe = MagicMock(spec=FeatureExtractionPipeline)
        pipeline_info = handle_transformers_pipeline(pipe)
        assert pipeline_info is not None
        assert pipeline_info["inputs"].label == "Input"
        assert pipeline_info["outputs"].label == "Output"

    def test_fill_mask_pipeline(self):
        pipe = MagicMock(spec=FillMaskPipeline)
        pipeline_info = handle_transformers_pipeline(pipe)
        assert pipeline_info is not None
        assert pipeline_info["inputs"].label == "Input"
        assert pipeline_info["outputs"].label == "Classification"

    def test_image_classification_pipeline(self):
        pipe = MagicMock(spec=ImageClassificationPipeline)
        pipeline_info = handle_transformers_pipeline(pipe)
        assert pipeline_info is not None
        assert pipeline_info["inputs"].label == "Input Image"
        assert pipeline_info["outputs"].label == "Classification"

    def test_question_answering_pipeline(self):
        pipe = MagicMock(spec=QuestionAnsweringPipeline)
        pipeline_info = handle_transformers_pipeline(pipe)
        assert pipeline_info is not None
        assert pipeline_info["inputs"][0].label == "Context"
        assert pipeline_info["inputs"][1].label == "Question"
        assert pipeline_info["outputs"][0].label == "Answer"
        assert pipeline_info["outputs"][1].label == "Score"

    def test_summarization_pipeline(self):
        pipe = MagicMock(spec=SummarizationPipeline)
        pipeline_info = handle_transformers_pipeline(pipe)
        assert pipeline_info is not None
        assert pipeline_info["inputs"].label == "Input"
        assert pipeline_info["outputs"].label == "Summary"

    def test_text_classification_pipeline(self):
        pipe = MagicMock(spec=TextClassificationPipeline)
        pipeline_info = handle_transformers_pipeline(pipe)
        assert pipeline_info is not None
        assert pipeline_info["inputs"].label == "Input"
        assert pipeline_info["outputs"].label == "Classification"

    def test_text_generation_pipeline(self):
        pipe = MagicMock(spec=TextGenerationPipeline)
        pipeline_info = handle_transformers_pipeline(pipe)
        assert pipeline_info is not None
        assert pipeline_info["inputs"].label == "Input"
        assert pipeline_info["outputs"].label == "Output"

    def test_translation_pipeline(self):
        pipe = MagicMock(spec=TranslationPipeline)
        pipeline_info = handle_transformers_pipeline(pipe)
        assert pipeline_info is not None
        assert pipeline_info["inputs"].label == "Input"
        assert pipeline_info["outputs"].label == "Translation"

    def test_text2text_generation_pipeline(self):
        pipe = MagicMock(spec=Text2TextGenerationPipeline)
        pipeline_info = handle_transformers_pipeline(pipe)
        assert pipeline_info is not None
        assert pipeline_info["inputs"].label == "Input"
        assert pipeline_info["outputs"].label == "Generated Text"

    def test_zero_shot_classification_pipeline(self):
        pipe = MagicMock(spec=ZeroShotClassificationPipeline)
        pipeline_info = handle_transformers_pipeline(pipe)
        assert pipeline_info is not None
        assert pipeline_info["inputs"][0].label == "Input"
        assert (
            pipeline_info["inputs"][1].label == "Possible class names (comma-separated)"
        )
        assert pipeline_info["inputs"][2].label == "Allow multiple true classes"
        assert pipeline_info["outputs"].label == "Classification"

    def test_document_question_answering_pipeline(self):
        pipe = MagicMock(spec=DocumentQuestionAnsweringPipeline)
        pipeline_info = handle_transformers_pipeline(pipe)
        assert pipeline_info is not None
        assert pipeline_info["inputs"][0].label == "Input Document"
        assert pipeline_info["inputs"][1].label == "Question"
        assert pipeline_info["outputs"].label == "Label"

    def test_visual_question_answering_pipeline(self):
        pipe = MagicMock(spec=VisualQuestionAnsweringPipeline)
        pipeline_info = handle_transformers_pipeline(pipe)
        assert pipeline_info is not None
        assert pipeline_info["inputs"][0].label == "Input Image"
        assert pipeline_info["inputs"][1].label == "Question"
        assert pipeline_info["outputs"].label == "Score"

    def test_image_to_text_pipeline(self):
        pipe = MagicMock(spec=ImageToTextPipeline)
        pipeline_info = handle_transformers_pipeline(pipe)
        assert pipeline_info is not None
        assert pipeline_info["inputs"].label == "Input Image"
        assert pipeline_info["outputs"].label == "Text"

    def test_unsupported_pipeline(self):
        pipe = MagicMock()
        with self.assertRaises(ValueError):
            handle_transformers_pipeline(pipe)


class TestHandleDiffusersPipelines(unittest.TestCase):
    def test_stable_diffusion_pipeline(self):
        pipe = MagicMock(spec=StableDiffusionPipeline)
        pipeline_info = handle_diffusers_pipeline(pipe)
        assert pipeline_info is not None
        assert pipeline_info["inputs"][0].label == "Prompt"
        assert pipeline_info["inputs"][1].label == "Negative prompt"
        assert pipeline_info["outputs"].label == "Generated Image"

    def test_stable_diffusion_img2img_pipeline(self):
        pipe = MagicMock(spec=StableDiffusionImg2ImgPipeline)
        pipeline_info = handle_diffusers_pipeline(pipe)
        assert pipeline_info is not None
        assert pipeline_info["inputs"][0].label == "Prompt"
        assert pipeline_info["inputs"][1].label == "Negative prompt"
        assert pipeline_info["outputs"].label == "Generated Image"

    def test_stable_diffusion_inpaint_pipeline(self):
        pipe = MagicMock(spec=StableDiffusionInpaintPipeline)
        pipeline_info = handle_diffusers_pipeline(pipe)
        assert pipeline_info is not None
        assert pipeline_info["inputs"][0].label == "Prompt"
        assert pipeline_info["inputs"][1].label == "Negative prompt"
        assert pipeline_info["outputs"].label == "Generated Image"

    def test_stable_diffusion_depth2img_pipeline(self):
        pipe = MagicMock(spec=StableDiffusionDepth2ImgPipeline)
        pipeline_info = handle_diffusers_pipeline(pipe)
        assert pipeline_info is not None
        assert pipeline_info["inputs"][0].label == "Prompt"
        assert pipeline_info["inputs"][1].label == "Negative prompt"
        assert pipeline_info["outputs"].label == "Generated Image"

    def test_stable_diffusion_image_variation_pipeline(self):
        pipe = MagicMock(spec=StableDiffusionImageVariationPipeline)
        pipeline_info = handle_diffusers_pipeline(pipe)
        assert pipeline_info is not None
        assert pipeline_info["inputs"][0].label == "Image"
        assert pipeline_info["outputs"].label == "Generated Image"

    def test_stable_diffusion_instruct_pix2pix_pipeline(self):
        pipe = MagicMock(spec=StableDiffusionInstructPix2PixPipeline)
        pipeline_info = handle_diffusers_pipeline(pipe)
        assert pipeline_info is not None
        assert pipeline_info["inputs"][0].label == "Prompt"
        assert pipeline_info["inputs"][1].label == "Negative prompt"
        assert pipeline_info["outputs"].label == "Generated Image"

    def test_stable_diffusion_upscale_pipeline(self):
        pipe = MagicMock(spec=StableDiffusionUpscalePipeline)
        pipeline_info = handle_diffusers_pipeline(pipe)
        assert pipeline_info is not None
        assert pipeline_info["inputs"][0].label == "Prompt"
        assert pipeline_info["inputs"][1].label == "Negative prompt"
        assert pipeline_info["outputs"].label == "Generated Image"

    def test_unsupported_pipeline(self):
        pipe = MagicMock()
        with self.assertRaises(ValueError):
            handle_transformers_pipeline(pipe)