test_model_zoo: adapt (load_models instead of load_model)

This commit is contained in:
Robert Sachunsky 2026-05-12 19:02:43 +02:00
parent 21ecb043f7
commit 1ed633bc25

View file

@ -6,11 +6,11 @@ def test_trocr1(
model_zoo = EynollahModelZoo(model_dir) model_zoo = EynollahModelZoo(model_dir)
try: try:
from transformers import TrOCRProcessor, VisionEncoderDecoderModel from transformers import TrOCRProcessor, VisionEncoderDecoderModel
model_zoo.load_model('trocr_processor') model_zoo.load_models('trocr_processor',
proc = model_zoo.get('trocr_processor', TrOCRProcessor) ('ocr', 'tr'))
proc = model_zoo.get('trocr_processor')
assert isinstance(proc, TrOCRProcessor) assert isinstance(proc, TrOCRProcessor)
model_zoo.load_model('ocr', 'tr') model = model_zoo.get('ocr')
model = model_zoo.get('ocr', VisionEncoderDecoderModel)
assert isinstance(model, VisionEncoderDecoderModel) assert isinstance(model, VisionEncoderDecoderModel)
except ImportError: except ImportError:
pass pass