diff --git a/dan/ocr/predict/prediction.py b/dan/ocr/predict/prediction.py
index 930e54c20dfeb41adb01b4d16d0f9a5ede1f301f..f88676af7e2b40c87aa09373a677563c1bb532a0 100644
--- a/dan/ocr/predict/prediction.py
+++ b/dan/ocr/predict/prediction.py
@@ -370,11 +370,10 @@ def process_batch(
 
         # Return LM results
         if use_language_model:
-            result["language_model"] = {}
-            result["language_model"]["text"] = prediction["language_model"]["text"][idx]
-            result["language_model"]["confidence"] = prediction["language_model"][
-                "confidence"
-            ][idx]
+            result["language_model"] = {
+                "text": prediction["language_model"]["text"][idx],
+                "confidence": prediction["language_model"]["confidence"][idx],
+            }
 
         # Return extracted objects (coordinates, text, confidence)
         if predict_objects:
diff --git a/dan/utils.py b/dan/utils.py
index fcd7e7afee411924f1bec8bc34037212c5de85d8..c65df263f1789e2799e1cc9c8cadf1b63dfea9f7 100644
--- a/dan/utils.py
+++ b/dan/utils.py
@@ -36,7 +36,7 @@ class LMTokenMapping(NamedTuple):
         return {a.display: a.encoded for a in self}
 
     def encode_token(self, token: str) -> str:
-        return self.encode[token] if token in self.encode else token
+        return self.encode.get(token, token)
 
 
 class EntityType(NamedTuple):