From cc82ca28600b89f2df86e027b4f37fab75880ff1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ga=C3=ABtan=20Caillaut?= Date: Mon, 18 Sep 2023 11:57:13 +0200 Subject: [PATCH] allow specifying model tag and batch size from environment variable --- app.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/app.py b/app.py index 0ddb62b..ac75199 100644 --- a/app.py +++ b/app.py @@ -23,8 +23,9 @@ from tqdm import tqdm SAVE_DIR = Path("./pdfs") -BATCHSIZE = os.environ.get("NOUGAT_BATCHSIZE", 6) -NOUGAT_CHECKPOINT = get_checkpoint() +BATCHSIZE = int(os.environ.get("NOUGAT_BATCHSIZE", 6)) +MODEL_TAG=os.environ.get("NOUGAT_MODEL_TAG", "0.1.0-small") +NOUGAT_CHECKPOINT = get_checkpoint(model_tag=MODEL_TAG) if NOUGAT_CHECKPOINT is None: print( "Set environment variable 'NOUGAT_CHECKPOINT' with a path to the model checkpoint!."