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

Add `NerModel.tokenizer`

parent 69c56bf1
Pipeline #147174 passed with stage
in 9 minutes and 12 seconds
......@@ -11,7 +11,7 @@ from adaptor.evaluators.token_classification import MeanFScore
from adaptor.objectives.MLM import MaskedLanguageModeling
from adaptor.lang_module import LangModule
from adaptor.utils import StoppingStrategy, AdaptationArguments
from transformers import AutoModelForTokenClassification
from transformers import AutoModelForTokenClassification, AutoTokenizer
from ..config import CONFIG as _CONFIG
from ..document import Document, Sentence
......@@ -42,6 +42,11 @@ class NerModel:
def __init__(self, model_name_or_basename: str):
self.model_name_or_basename = model_name_or_basename
def tokenizer(self) -> AutoTokenizer:
tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_basename)
return tokenizer
def model(self) -> AutoModelForTokenClassification:
model = AutoModelForTokenClassification.from_pretrained(self.model_name_or_basename)
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