qurator-spk/eynollah
1
0
Fork 0
mirror of https://github.com/qurator-spk/eynollah.git synced 2025-11-10 06:34:11 +01:00
Code Issues Projects Releases Wiki Activity

model_zoo: make type str to reduce importing overhead

Browse source
This commit is contained in:
kba 2025-10-29 19:08:32 +01:00
parent a913bdf7dc
commit 5e22e9db64
6 changed files with 63 additions and 59 deletions
Download patch file Download diff file Expand all files Collapse all files

33
tests/test_model_zoo.py
View file

@ -1,19 +1,16 @@
from pathlib import Path
from eynollah.model_zoo import EynollahModelZoo
from eynollah.model_zoo import EynollahModelZoo, TrOCRProcessor, VisionEncoderDecoderModel
testdir = Path(__file__).parent.resolve()
MODELS_DIR = testdir.parent
def test_trocr1():
model_zoo = EynollahModelZoo(str(MODELS_DIR))
model_zoo.load_model('trocr_processor')
proc = model_zoo.get('trocr_processor', TrOCRProcessor)
assert isinstance(proc, TrOCRProcessor)
model_zoo.load_model('ocr', 'tr')
model = model_zoo.get('ocr')
assert isinstance(model, VisionEncoderDecoderModel)
print(proc)
test_trocr1()
def test_trocr1(
model_dir,
):
model_zoo = EynollahModelZoo(model_dir)
try:
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
model_zoo.load_model('trocr_processor')
proc = model_zoo.get('trocr_processor', TrOCRProcessor)
assert isinstance(proc, TrOCRProcessor)
model_zoo.load_model('ocr', 'tr')
model = model_zoo.get('ocr', VisionEncoderDecoderModel)
assert isinstance(model, VisionEncoderDecoderModel)
except ImportError:
pass