eynollah/tests/test_model_zoo.py

from eynollah.model_zoo import EynollahModelZoo

def test_trocr1(
    model_dir,
):
    model_zoo = EynollahModelZoo(model_dir)
    try:
        from transformers import TrOCRProcessor, VisionEncoderDecoderModel
        model_zoo.load_models('trocr_processor')
        proc = model_zoo.get('trocr_processor')
        assert isinstance(proc, TrOCRProcessor)
        model_zoo.load_models(['ocr', 'tr'])
        model = model_zoo.get('ocr')
        assert isinstance(model, VisionEncoderDecoderModel)
    except ImportError:
        pass