Fix tests

tests/test_prediction.py

@@ -3,6 +3,7 @@
 import json
 import shutil
 
+import numpy as np
 import pytest
 import yaml
 
@@ -545,22 +546,30 @@ def test_run_prediction_batch(
                 {
                     "text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241",
                     "language_model": {
-                        "text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241"
+                        "text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241",
+                        "confidence": 0.9226743371961854,
                     },
                 },
                 {
                     "text": "ⓈTemplié ⒻMarcelle Ⓑ93 ⓁS Ⓚch ⓄE dactylo Ⓟ18376",
                     "language_model": {
-                        "text": "ⓈTemplié ⒻMarcelle Ⓑ93 ⓁS Ⓚch ⓄE dactylo Ⓟ18376"
+                        "text": "ⓈTemplié ⒻMarcelle Ⓑ93 ⓁS Ⓚch ⓄE dactylo Ⓟ18376",
+                        "confidence": 0.8759829104754289,
                     },
                 },
                 {
                     "text": "Ⓢd ⒻCharles Ⓑ11 ⓁP ⒸC ⓀF Ⓞd Ⓟ14 31",
-                    "language_model": {"text": "Ⓢd ⒻCharles Ⓑ11 ⓁP ⒸC ⓀF Ⓞd Ⓟ1431"},
+                    "language_model": {
+                        "text": "Ⓢd ⒻCharles Ⓑ11 ⓁP ⒸC ⓀF Ⓞd Ⓟ14 31",
+                        "confidence": 0.864021797502254,
+                    },
                 },
                 {
                     "text": "ⓈNaudin ⒻMarie Ⓑ53 ⓁS Ⓒv ⓀBelle mère",
-                    "language_model": {"text": "ⓈNaudin ⒻMarie Ⓑ53 ⓁS Ⓒv ⓀBelle mère"},
+                    "language_model": {
+                        "text": "ⓈNaudin ⒻMarie Ⓑ53 ⓁS Ⓒv ⓀBelle mère",
+                        "confidence": 0.8903665579889012,
+                    },
                 },
             ],
         ),
@@ -576,22 +585,30 @@ def test_run_prediction_batch(
                 {
                     "text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241",
                     "language_model": {
-                        "text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241"
+                        "text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241",
+                        "confidence": 0.8982517863786614,
                     },
                 },
                 {
                     "text": "ⓈTemplié ⒻMarcelle Ⓑ93 ⓁS Ⓚch ⓄE dactylo Ⓟ18376",
                     "language_model": {
-                        "text": "ⓈTemplié ⒻMarcelle Ⓑ93 ⓁS Ⓚch ⓄE dactylo Ⓟ18376"
+                        "text": "ⓈTemplié ⒻMarcelle Ⓑ93 ⓁS Ⓚch ⓄE dactylo Ⓟ18376",
+                        "confidence": 0.8386571587822831,
                     },
                 },
                 {
                     "text": "Ⓢd ⒻCharles Ⓑ11 ⓁP ⒸC ⓀF Ⓞd Ⓟ14 31",
-                    "language_model": {"text": "Ⓢd ⒻCharles Ⓑ11 ⓁP ⒸC ⓀF Ⓞd Ⓟ14331"},
+                    "language_model": {
+                        "text": "Ⓢd ⒻCharles Ⓑ11 ⓁP ⒸC ⓀF Ⓞd Ⓟ14331",
+                        "confidence": 0.8334836549049839,
+                    },
                 },
                 {
                     "text": "ⓈNaudin ⒻMarie Ⓑ53 ⓁS Ⓒv ⓀBelle mère",
-                    "language_model": {"text": "ⓈNaudin ⒻMarie Ⓑ53 ⓁS Ⓒv ⓀBelle mère"},
+                    "language_model": {
+                        "text": "ⓈNaudin ⒻMarie Ⓑ53 ⓁS Ⓒv ⓀBelle mère",
+                        "confidence": 0.8565623750166133,
+                    },
                 },
             ],
         ),
@@ -612,12 +629,10 @@ def test_run_prediction_batch(
         ),
     ),
 )
-@pytest.mark.parametrize("batch_size", [1, 2])
 def test_run_prediction_language_model(
     image_names,
     language_model_weight,
     expected_predictions,
-    batch_size,
     tmp_path,
 ):
     # Make tmpdir and copy needed images inside
@@ -655,7 +670,7 @@ def test_run_prediction_language_model(
         max_object_height=None,
         image_extension=".png",
         gpu_device=None,
-        batch_size=batch_size,
+        batch_size=1,
         tokens=parse_tokens(PREDICTION_DATA_PATH / "tokens.yml"),
         start_token=None,
         use_language_model=True,
@@ -668,7 +683,15 @@ def test_run_prediction_language_model(
         assert prediction["text"] == expected_prediction["text"]
 
         if language_model_weight > 0:
+            print(
+                prediction["language_model"]["text"],
+                prediction["language_model"]["confidence"],
+            )
             assert (
                 prediction["language_model"]["text"]
                 == expected_prediction["language_model"]["text"]
             )
+            assert np.isclose(
+                prediction["language_model"]["confidence"],
+                expected_prediction["language_model"]["confidence"],
+            )