fix: one extra token prediction #104

Kohulan · Jun 11, 2024 · 6e0e483 · 6e0e483
1 parent 52635a4
commit 6e0e483
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 6 deletions.
diff --git a/DECIMER/decimer.py b/DECIMER/decimer.py
@@ -117,7 +117,7 @@ def detokenize_output_add_confidence(
     decoded_prediction_with_confidence = list(
         [(utils.decoder(tok), conf) for tok, conf in prediction_with_confidence_]
     )
-    decoded_prediction_with_confidence.append(prediction_with_confidence_[-1])
+
     return decoded_prediction_with_confidence
 
 

diff --git a/DECIMER/utils.py b/DECIMER/utils.py
@@ -83,11 +83,10 @@ def ensure_models(default_path: str, model_urls: dict) -> dict:
     }
     for model_name, model_url in model_urls.items():
         model_path = os.path.join(default_path, f"{model_name}_model")
-        if (
-                os.path.exists(model_path)
-                and os.stat(os.path.join(model_path, "saved_model.pb")).st_size != model_sizes.get(model_name)
-        ):
-            print(f'Working with model {model_name}')
+        if os.path.exists(model_path) and os.stat(
+            os.path.join(model_path, "saved_model.pb")
+        ).st_size != model_sizes.get(model_name):
+            print(f"Working with model {model_name}")
             shutil.rmtree(model_path)
             config.download_trained_weights(model_url, default_path)
         elif not os.path.exists(model_path):