wip

tests/test_model_zoo.py

@@ -0,0 +1,19 @@
+from pathlib import Path
+
+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()