Commit 7263450d authored by Vít Novotný's avatar Vít Novotný
Browse files

Make `NerModel.__init__()` lazy and add `str()` and `repr()`

parent b78567bb
......@@ -36,8 +36,18 @@ class NerModel:
NUM_VALIDATION_SAMPLES = CONFIG.getint('number_of_validation_samples')
STOPPING_PATIENCE = CONFIG.getint('stopping_patience')
def __init__(self, model: AutoModelForTokenClassification):
self.model = model
def __init__(self, model_name_or_basename: str):
self.model_name_or_basename = model_name_or_basename
def model(self) -> AutoModelForTokenClassification:
model = AutoModelForTokenClassification.from_pretrained(self.model_name_or_basename)
def __str__(self) -> str:
return self.model_name_or_basename
def __repr__(self) -> str:
return '{}: {}'.format(self.__class__.__name__, self)
def train_and_save(cls, model_checkpoint_basename: str, model_basename: str,
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment