test/tasks/text_gradio_test.py

from pathlib import Path

from deep_learning.data.spec import TokenizerBundle
from deep_learning.pipeline.impl.text_generation.pipeline import TextInferenceBundle
from deep_learning.ui.gradio.text_app import TextGenerationAppBuilder
from deep_learning.ui.gradio.text_deployment import (
    load_poetry_gpt_deployment_inference_artifact,
    load_poetry_rnn_deployment_inference_artifact,
    load_wiki_gpt_deployment_inference_artifact
)


class DummyModel:
    pass


class DummyTokenizer:
    def vocabulary_size(self):
        return 32


class DummyInput:
    def __init__(self, hidden_dim: int):
        self.shape = (None, hidden_dim)


class DummyStatefulModel:
    def __init__(self):
        self.inputs = [
            DummyInput(1),
            DummyInput(16),
            DummyInput(16),
            DummyInput(16),
            DummyInput(16)
        ]


class DummyCheckpointLoadRules:
    def __init__(self, checkpoint_dir: Path):
        self.checkpoint_dir = checkpoint_dir

    def resolve_test_rule(self, default_dirs):
        return {
            "dirs": [self.checkpoint_dir]
        }


class DummyPipeline:
    def __init__(self, checkpoint_dir: Path):
        self.name = "test_text"
        self.checkpoint_dir = checkpoint_dir
        self.checkpoint_load_rules = DummyCheckpointLoadRules(checkpoint_dir)


def _create_text_app_builder(tmp_path, load_inference_artifact, load_inference_resource=None):
    checkpoint_dir = tmp_path / "checkpoints"
    checkpoint_dir.mkdir()
    (checkpoint_dir / "model_epoch_001.keras").write_text("stub", encoding="utf-8")
    return TextGenerationAppBuilder(
        pipeline=DummyPipeline(checkpoint_dir),
        title="测试文本页面",
        load_inference_artifact=load_inference_artifact,
        load_inference_resource=load_inference_resource
    )


def _dummy_text_inference_bundle() -> TextInferenceBundle:
    return TextInferenceBundle(
        tokenizer_bundle=TokenizerBundle(
            tokenizer=lambda text: [],
            decode=lambda token_ids: "",
            end_of_text=99,
            vocab_size=32,
            vocab_path="saved/vocab/test.txt"
        ),
        docs_ds=None,
        max_length=16,
        sample_fn=None
    )


def test_text_app_builder_model_info_does_not_load_model(tmp_path):
    """验证文本页面展示模型信息时只读取推理资源，不加载完整模型。"""
    builder = _create_text_app_builder(
        tmp_path,
        lambda checkpoint_rule: (_ for _ in ()).throw(AssertionError("模型信息不应加载模型")),
        load_inference_resource=_dummy_text_inference_bundle
    )

    model_info = builder.get_model_info()

    assert "model_epoch_001.keras" in model_info
    assert "saved/vocab/test.txt" in model_info


def test_wiki_gpt_app_builder_loads_deployment_resource_without_training_data(tmp_path, monkeypatch):
    """验证 Wiki GPT 页面加载部署模型和仓库词表时不会读取训练数据。"""
    builder = _create_text_app_builder(tmp_path, load_wiki_gpt_deployment_inference_artifact)
    monkeypatch.setattr(
        "deep_learning.models.mini_gpt.model_builder.keras.models.load_model",
        lambda path, custom_objects: DummyModel()
    )
    monkeypatch.setattr(
        "deep_learning.ui.gradio.text_deployment.sentence_piece",
        lambda: (DummyTokenizer(), 99, lambda token_ids: "")
    )

    inference_artifact, resource = builder._load_inference_artifact()

    assert isinstance(inference_artifact.model, DummyModel)
    assert isinstance(resource, TextInferenceBundle)
    assert resource.tokenizer_bundle.vocab_path.endswith("saved/vocab/sentencepiece/vocabulary.proto")


def test_poetry_gpt_app_builder_loads_deployment_resource_without_training_data(tmp_path, monkeypatch):
    """验证 Poetry GPT 页面加载部署模型和仓库词表时不会读取训练数据。"""
    builder = _create_text_app_builder(tmp_path, load_poetry_gpt_deployment_inference_artifact)
    monkeypatch.setattr(
        "deep_learning.models.mini_gpt.model_builder.keras.models.load_model",
        lambda path, custom_objects: DummyModel()
    )

    inference_artifact, resource = builder._load_inference_artifact()

    assert isinstance(inference_artifact.model, DummyModel)
    assert isinstance(resource, TextInferenceBundle)
    assert resource.tokenizer_bundle.vocab_path.endswith("saved/vocab/poetry/vocab.txt")


def test_poetry_rnn_app_builder_loads_deployment_resource_without_training_data(tmp_path, monkeypatch):
    """验证 Poetry RNN 页面加载部署模型和仓库词表时不会读取训练数据。"""
    builder = _create_text_app_builder(tmp_path, load_poetry_rnn_deployment_inference_artifact)
    monkeypatch.setattr(
        "deep_learning.models.rnn.keras.models.load_model",
        lambda path: DummyStatefulModel()
    )

    inference_artifact, resource = builder._load_inference_artifact()

    assert isinstance(inference_artifact.model, DummyStatefulModel)
    assert isinstance(resource, TextInferenceBundle)
    assert resource.tokenizer_bundle.vocab_path.endswith("saved/vocab/poetry/vocab.txt")