eynollah/tests/test_model_zoo.py
2026-05-11 12:22:24 +02:00

16 lines
536 B
Python

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