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):