From 8b9662480abd6104869c65f52b4edc80c5bf635e Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Sun, 10 Oct 2021 18:08:21 +0200 Subject: Update transformer models --- text_recognizer/models/transformer.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) (limited to 'text_recognizer/models/transformer.py') diff --git a/text_recognizer/models/transformer.py b/text_recognizer/models/transformer.py index dbdc7f2..0acc30a 100644 --- a/text_recognizer/models/transformer.py +++ b/text_recognizer/models/transformer.py @@ -1,10 +1,11 @@ """PyTorch Lightning model for base Transformers.""" -from typing import Tuple, Set +from typing import Tuple, Type, Set import attr import torch from torch import Tensor +from text_recognizer.data.base_mapping import AbstractMapping from text_recognizer.models.metrics import CharacterErrorRate from text_recognizer.models.base import BaseLitModel @@ -13,6 +14,8 @@ from text_recognizer.models.base import BaseLitModel class TransformerLitModel(BaseLitModel): """A PyTorch Lightning model for transformer networks.""" + mapping: Type[AbstractMapping] = attr.ib() + max_output_len: int = attr.ib(default=451) start_token: str = attr.ib(default="") end_token: str = attr.ib(default="") -- cgit v1.2.3-70-g09d2