rasa-pro 3.8.18__py3-none-any.whl → 3.9.15__py3-none-any.whl

rasa/nlu/classifiers/diet_classifier.py

@@ -1,17 +1,18 @@
 from __future__ import annotations
-
 import copy
 import logging
 from collections import defaultdict
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Text, Tuple, Union, TypeVar, Type
+
+from rasa.exceptions import ModelNotFound
+from rasa.nlu.featurizers.featurizer import Featurizer
 
 import numpy as np
 import scipy.sparse
 import tensorflow as tf
 
-from rasa.exceptions import ModelNotFound
-from rasa.nlu.featurizers.featurizer import Featurizer
+from typing import Any, Dict, List, Optional, Text, Tuple, Union, TypeVar, Type
+
 from rasa.engine.graph import ExecutionContext, GraphComponent
 from rasa.engine.recipes.default_recipe import DefaultV1Recipe
 from rasa.engine.storage.resource import Resource
@@ -19,21 +20,18 @@ from rasa.engine.storage.storage import ModelStorage
 from rasa.nlu.extractors.extractor import EntityExtractorMixin
 from rasa.nlu.classifiers.classifier import IntentClassifier
 import rasa.shared.utils.io
+import rasa.utils.io as io_utils
 import rasa.nlu.utils.bilou_utils as bilou_utils
 from rasa.shared.constants import DIAGNOSTIC_DATA
 from rasa.nlu.extractors.extractor import EntityTagSpec
 from rasa.nlu.classifiers import LABEL_RANKING_LENGTH
 from rasa.utils import train_utils
 from rasa.utils.tensorflow import rasa_layers
-from rasa.utils.tensorflow.feature_array import (
-    FeatureArray,
-    serialize_nested_feature_arrays,
-    deserialize_nested_feature_arrays,
-)
 from rasa.utils.tensorflow.models import RasaModel, TransformerRasaModel
 from rasa.utils.tensorflow.model_data import (
     RasaModelData,
     FeatureSignature,
+    FeatureArray,
 )
 from rasa.nlu.constants import TOKENS_NAMES, DEFAULT_TRANSFORMER_SIZE
 from rasa.shared.nlu.constants import (
@@ -120,6 +118,7 @@ LABEL_SUB_KEY = IDS
 
 POSSIBLE_TAGS = [ENTITY_ATTRIBUTE_TYPE, ENTITY_ATTRIBUTE_ROLE, ENTITY_ATTRIBUTE_GROUP]
 
+
 DIETClassifierT = TypeVar("DIETClassifierT", bound="DIETClassifier")
 
 
@@ -511,7 +510,6 @@ class DIETClassifier(GraphComponent, IntentClassifier, EntityExtractorMixin):
     def _extract_features(
         self, message: Message, attribute: Text
     ) -> Dict[Text, Union[scipy.sparse.spmatrix, np.ndarray]]:
-
         (
             sparse_sequence_features,
             sparse_sentence_features,
@@ -781,7 +779,6 @@ class DIETClassifier(GraphComponent, IntentClassifier, EntityExtractorMixin):
         sparse_feature_sizes: Dict[Text, Dict[Text, List[int]]],
         label_attribute: Optional[Text] = None,
     ) -> Dict[Text, Dict[Text, List[int]]]:
-
         if label_attribute in sparse_feature_sizes:
             del sparse_feature_sizes[label_attribute]
         return sparse_feature_sizes
@@ -1086,24 +1083,18 @@ class DIETClassifier(GraphComponent, IntentClassifier, EntityExtractorMixin):
 
             self.model.save(str(tf_model_file))
 
-            # save data example
-            serialize_nested_feature_arrays(
-                self._data_example,
-                model_path / f"{file_name}.data_example.st",
-                model_path / f"{file_name}.data_example_metadata.json",
+            io_utils.pickle_dump(
+                model_path / f"{file_name}.data_example.pkl", self._data_example
             )
-            # save label data
-            serialize_nested_feature_arrays(
-                dict(self._label_data.data) if self._label_data is not None else {},
-                model_path / f"{file_name}.label_data.st",
-                model_path / f"{file_name}.label_data_metadata.json",
-            )
-
-            rasa.shared.utils.io.dump_obj_as_json_to_file(
-                model_path / f"{file_name}.sparse_feature_sizes.json",
+            io_utils.pickle_dump(
+                model_path / f"{file_name}.sparse_feature_sizes.pkl",
                 self._sparse_feature_sizes,
             )
-            rasa.shared.utils.io.dump_obj_as_json_to_file(
+            io_utils.pickle_dump(
+                model_path / f"{file_name}.label_data.pkl",
+                dict(self._label_data.data) if self._label_data is not None else {},
+            )
+            io_utils.json_pickle(
                 model_path / f"{file_name}.index_label_id_mapping.json",
                 self.index_label_id_mapping,
             )
@@ -1192,22 +1183,15 @@ class DIETClassifier(GraphComponent, IntentClassifier, EntityExtractorMixin):
     ]:
         file_name = cls.__name__
 
-        # load data example
-        data_example = deserialize_nested_feature_arrays(
-            str(model_path / f"{file_name}.data_example.st"),
-            str(model_path / f"{file_name}.data_example_metadata.json"),
+        data_example = io_utils.pickle_load(
+            model_path / f"{file_name}.data_example.pkl"
         )
-        # load label data
-        loaded_label_data = deserialize_nested_feature_arrays(
-            str(model_path / f"{file_name}.label_data.st"),
-            str(model_path / f"{file_name}.label_data_metadata.json"),
-        )
-        label_data = RasaModelData(data=loaded_label_data)
-
-        sparse_feature_sizes = rasa.shared.utils.io.read_json_file(
-            model_path / f"{file_name}.sparse_feature_sizes.json"
+        label_data = io_utils.pickle_load(model_path / f"{file_name}.label_data.pkl")
+        label_data = RasaModelData(data=label_data)
+        sparse_feature_sizes = io_utils.pickle_load(
+            model_path / f"{file_name}.sparse_feature_sizes.pkl"
         )
-        index_label_id_mapping = rasa.shared.utils.io.read_json_file(
+        index_label_id_mapping = io_utils.json_unpickle(
             model_path / f"{file_name}.index_label_id_mapping.json"
         )
         entity_tag_specs = rasa.shared.utils.io.read_json_file(
@@ -1227,6 +1211,7 @@ class DIETClassifier(GraphComponent, IntentClassifier, EntityExtractorMixin):
             for tag_spec in entity_tag_specs
         ]
 
+        # jsonpickle converts dictionary keys to strings
         index_label_id_mapping = {
             int(key): value for key, value in index_label_id_mapping.items()
         }
@@ -1280,7 +1265,6 @@ class DIETClassifier(GraphComponent, IntentClassifier, EntityExtractorMixin):
         config: Dict[Text, Any],
         finetune_mode: bool,
     ) -> "RasaModel":
-
         predict_data_example = RasaModelData(
             label_key=model_data_example.label_key,
             data={
@@ -1467,10 +1451,10 @@ class DIET(TransformerRasaModel):
         # everything using a transformer and optionally also do masked language
         # modeling.
         self.text_name = TEXT
-        self._tf_layers[
-            f"sequence_layer.{self.text_name}"
-        ] = rasa_layers.RasaSequenceLayer(
-            self.text_name, self.data_signature[self.text_name], self.config
+        self._tf_layers[f"sequence_layer.{self.text_name}"] = (
+            rasa_layers.RasaSequenceLayer(
+                self.text_name, self.data_signature[self.text_name], self.config
+            )
         )
         if self.config[MASKED_LM]:
             self._prepare_mask_lm_loss(self.text_name)
@@ -1488,10 +1472,10 @@ class DIET(TransformerRasaModel):
                 {SPARSE_INPUT_DROPOUT: False, DENSE_INPUT_DROPOUT: False}
             )
 
-            self._tf_layers[
-                f"feature_combining_layer.{self.label_name}"
-            ] = rasa_layers.RasaFeatureCombiningLayer(
-                self.label_name, self.label_signature[self.label_name], label_config
+            self._tf_layers[f"feature_combining_layer.{self.label_name}"] = (
+                rasa_layers.RasaFeatureCombiningLayer(
+                    self.label_name, self.label_signature[self.label_name], label_config
+                )
             )
 
             self._prepare_ffnn_layer(
@@ -1523,7 +1507,6 @@ class DIET(TransformerRasaModel):
         sequence_feature_lengths: tf.Tensor,
         name: Text,
     ) -> tf.Tensor:
-
         x, _ = self._tf_layers[f"feature_combining_layer.{name}"](
             (sequence_features, sentence_features, sequence_feature_lengths),
             training=self._training,
@@ -1705,7 +1688,6 @@ class DIET(TransformerRasaModel):
         return loss
 
     def _update_label_metrics(self, loss: tf.Tensor, acc: tf.Tensor) -> None:
-
         self.intent_loss.update_state(loss)
         self.intent_acc.update_state(acc)
 
@@ -1864,7 +1846,6 @@ class DIET(TransformerRasaModel):
         combined_sequence_sentence_feature_lengths: tf.Tensor,
         text_transformed: tf.Tensor,
     ) -> Dict[Text, tf.Tensor]:
-
         if self.all_labels_embed is None:
             raise ValueError(
                 "The model was not prepared for prediction. "

rasa/nlu/classifiers/logistic_regression_classifier.py

@@ -1,21 +1,22 @@
 from typing import Any, Text, Dict, List, Type, Tuple
 
+import joblib
 import structlog
 from scipy.sparse import hstack, vstack, csr_matrix
 from sklearn.exceptions import NotFittedError
 from sklearn.linear_model import LogisticRegression
 from sklearn.utils.validation import check_is_fitted
 
-from rasa.engine.graph import ExecutionContext, GraphComponent
-from rasa.engine.recipes.default_recipe import DefaultV1Recipe
 from rasa.engine.storage.resource import Resource
 from rasa.engine.storage.storage import ModelStorage
+from rasa.engine.recipes.default_recipe import DefaultV1Recipe
+from rasa.engine.graph import ExecutionContext, GraphComponent
 from rasa.nlu.classifiers import LABEL_RANKING_LENGTH
-from rasa.nlu.classifiers.classifier import IntentClassifier
 from rasa.nlu.featurizers.featurizer import Featurizer
-from rasa.shared.nlu.constants import TEXT, INTENT
-from rasa.shared.nlu.training_data.message import Message
+from rasa.nlu.classifiers.classifier import IntentClassifier
 from rasa.shared.nlu.training_data.training_data import TrainingData
+from rasa.shared.nlu.training_data.message import Message
+from rasa.shared.nlu.constants import TEXT, INTENT
 from rasa.utils.tensorflow.constants import RANKING_LENGTH
 
 structlogger = structlog.get_logger()
@@ -183,11 +184,9 @@ class LogisticRegressionClassifier(IntentClassifier, GraphComponent):
 
     def persist(self) -> None:
         """Persist this model into the passed directory."""
-        import skops.io as sio
-
         with self._model_storage.write_to(self._resource) as model_dir:
-            path = model_dir / f"{self._resource.name}.skops"
-            sio.dump(self.clf, path)
+            path = model_dir / f"{self._resource.name}.joblib"
+            joblib.dump(self.clf, path)
             structlogger.debug(
                 "logistic_regression_classifier.persist",
                 event_info=f"Saved intent classifier to '{path}'.",
@@ -203,21 +202,9 @@ class LogisticRegressionClassifier(IntentClassifier, GraphComponent):
         **kwargs: Any,
     ) -> "LogisticRegressionClassifier":
         """Loads trained component (see parent class for full docstring)."""
-        import skops.io as sio
-
         try:
             with model_storage.read_from(resource) as model_dir:
-                classifier_file = model_dir / f"{resource.name}.skops"
-                unknown_types = sio.get_untrusted_types(file=classifier_file)
-
-                if unknown_types:
-                    structlogger.error(
-                        f"Untrusted types found when loading {classifier_file}!",
-                        unknown_types=unknown_types,
-                    )
-                    raise ValueError()
-
-                classifier = sio.load(classifier_file, trusted=unknown_types)
+                classifier = joblib.load(model_dir / f"{resource.name}.joblib")
                 component = cls(
                     config, execution_context.node_name, model_storage, resource
                 )

rasa/nlu/classifiers/sklearn_intent_classifier.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
-
 import logging
+from rasa.nlu.featurizers.dense_featurizer.dense_featurizer import DenseFeaturizer
 import typing
 import warnings
 from typing import Any, Dict, List, Optional, Text, Tuple, Type
@@ -8,18 +8,18 @@ from typing import Any, Dict, List, Optional, Text, Tuple, Type
 import numpy as np
 
 import rasa.shared.utils.io
+import rasa.utils.io as io_utils
 from rasa.engine.graph import GraphComponent, ExecutionContext
 from rasa.engine.recipes.default_recipe import DefaultV1Recipe
 from rasa.engine.storage.resource import Resource
 from rasa.engine.storage.storage import ModelStorage
-from rasa.nlu.classifiers import LABEL_RANKING_LENGTH
-from rasa.nlu.classifiers.classifier import IntentClassifier
-from rasa.nlu.featurizers.dense_featurizer.dense_featurizer import DenseFeaturizer
 from rasa.shared.constants import DOCS_URL_TRAINING_DATA_NLU
+from rasa.nlu.classifiers import LABEL_RANKING_LENGTH
 from rasa.shared.exceptions import RasaException
 from rasa.shared.nlu.constants import TEXT
-from rasa.shared.nlu.training_data.message import Message
+from rasa.nlu.classifiers.classifier import IntentClassifier
 from rasa.shared.nlu.training_data.training_data import TrainingData
+from rasa.shared.nlu.training_data.message import Message
 from rasa.utils.tensorflow.constants import FEATURIZERS
 
 logger = logging.getLogger(__name__)
@@ -266,20 +266,14 @@ class SklearnIntentClassifier(GraphComponent, IntentClassifier):
 
     def persist(self) -> None:
         """Persist this model into the passed directory."""
-        import skops.io as sio
-
         with self._model_storage.write_to(self._resource) as model_dir:
             file_name = self.__class__.__name__
-            classifier_file_name = model_dir / f"{file_name}_classifier.skops"
-            encoder_file_name = model_dir / f"{file_name}_encoder.json"
+            classifier_file_name = model_dir / f"{file_name}_classifier.pkl"
+            encoder_file_name = model_dir / f"{file_name}_encoder.pkl"
 
             if self.clf and self.le:
-                # convert self.le.classes_ (numpy array of strings) to a list in order
-                # to use json dump
-                rasa.shared.utils.io.dump_obj_as_json_to_file(
-                    encoder_file_name, list(self.le.classes_)
-                )
-                sio.dump(self.clf.best_estimator_, classifier_file_name)
+                io_utils.json_pickle(encoder_file_name, self.le.classes_)
+                io_utils.json_pickle(classifier_file_name, self.clf.best_estimator_)
 
     @classmethod
     def load(
@@ -292,36 +286,21 @@ class SklearnIntentClassifier(GraphComponent, IntentClassifier):
     ) -> SklearnIntentClassifier:
         """Loads trained component (see parent class for full docstring)."""
         from sklearn.preprocessing import LabelEncoder
-        import skops.io as sio
 
         try:
             with model_storage.read_from(resource) as model_dir:
                 file_name = cls.__name__
-                classifier_file = model_dir / f"{file_name}_classifier.skops"
+                classifier_file = model_dir / f"{file_name}_classifier.pkl"
 
                 if classifier_file.exists():
-                    unknown_types = sio.get_untrusted_types(file=classifier_file)
-
-                    if unknown_types:
-                        logger.error(
-                            f"Untrusted types ({unknown_types}) found when "
-                            f"loading {classifier_file}!"
-                        )
-                        raise ValueError()
-                    else:
-                        classifier = sio.load(classifier_file, trusted=unknown_types)
-
-                    encoder_file = model_dir / f"{file_name}_encoder.json"
-                    classes = rasa.shared.utils.io.read_json_file(encoder_file)
+                    classifier = io_utils.json_unpickle(classifier_file)
 
+                    encoder_file = model_dir / f"{file_name}_encoder.pkl"
+                    classes = io_utils.json_unpickle(encoder_file)
                     encoder = LabelEncoder()
-                    intent_classifier = cls(
-                        config, model_storage, resource, classifier, encoder
-                    )
-                    # convert list of strings (class labels) back to numpy array of
-                    # strings
-                    intent_classifier.transform_labels_str2num(classes)
-                    return intent_classifier
+                    encoder.classes_ = classes
+
+                    return cls(config, model_storage, resource, classifier, encoder)
         except ValueError:
             logger.debug(
                 f"Failed to load '{cls.__name__}' from model storage. Resource "

rasa/nlu/extractors/crf_entity_extractor.py

@@ -4,9 +4,9 @@ from collections import OrderedDict
 from enum import Enum
 import logging
 import typing
-from typing import Any, Dict, List, Optional, Text, Tuple, Callable, Type
 
 import numpy as np
+from typing import Any, Dict, List, Optional, Text, Tuple, Callable, Type
 
 import rasa.nlu.utils.bilou_utils as bilou_utils
 import rasa.shared.utils.io
@@ -41,9 +41,6 @@ if typing.TYPE_CHECKING:
     from sklearn_crfsuite import CRF
 
 
-CONFIG_FEATURES = "features"
-
-
 class CRFToken:
     def __init__(
         self,
@@ -63,29 +60,6 @@ class CRFToken:
         self.entity_role_tag = entity_role_tag
         self.entity_group_tag = entity_group_tag
 
-    def to_dict(self) -> Dict[str, Any]:
-        return {
-            "text": self.text,
-            "pos_tag": self.pos_tag,
-            "pattern": self.pattern,
-            "dense_features": [str(x) for x in list(self.dense_features)],
-            "entity_tag": self.entity_tag,
-            "entity_role_tag": self.entity_role_tag,
-            "entity_group_tag": self.entity_group_tag,
-        }
-
-    @classmethod
-    def create_from_dict(cls, data: Dict[str, Any]) -> "CRFToken":
-        return cls(
-            data["text"],
-            data["pos_tag"],
-            data["pattern"],
-            np.array([float(x) for x in data["dense_features"]]),
-            data["entity_tag"],
-            data["entity_role_tag"],
-            data["entity_group_tag"],
-        )
-
 
 class CRFEntityExtractorOptions(str, Enum):
     """Features that can be used for the 'CRFEntityExtractor'."""
@@ -114,6 +88,8 @@ class CRFEntityExtractorOptions(str, Enum):
 class CRFEntityExtractor(GraphComponent, EntityExtractorMixin):
     """Implements conditional random fields (CRF) to do named entity recognition."""
 
+    CONFIG_FEATURES = "features"
+
     function_dict: Dict[Text, Callable[[CRFToken], Any]] = {  # noqa: RUF012
         CRFEntityExtractorOptions.LOW: lambda crf_token: crf_token.text.lower(),
         CRFEntityExtractorOptions.TITLE: lambda crf_token: crf_token.text.istitle(),
@@ -132,7 +108,7 @@ class CRFEntityExtractor(GraphComponent, EntityExtractorMixin):
         CRFEntityExtractorOptions.DIGIT: lambda crf_token: crf_token.text.isdigit(),
         CRFEntityExtractorOptions.PATTERN: lambda crf_token: crf_token.pattern,
         CRFEntityExtractorOptions.TEXT_DENSE_FEATURES: (
-            lambda crf_token: CRFEntityExtractor._convert_dense_features_for_crfsuite(  # noqa: E501
+            lambda crf_token: CRFEntityExtractor._convert_dense_features_for_crfsuite(
                 crf_token
             )
         ),
@@ -161,7 +137,7 @@ class CRFEntityExtractor(GraphComponent, EntityExtractorMixin):
             # "is the preceding token in title case?"
             # POS features require SpacyTokenizer
             # pattern feature require RegexFeaturizer
-            CONFIG_FEATURES: [
+            CRFEntityExtractor.CONFIG_FEATURES: [
                 [
                     CRFEntityExtractorOptions.LOW,
                     CRFEntityExtractorOptions.TITLE,
@@ -224,7 +200,7 @@ class CRFEntityExtractor(GraphComponent, EntityExtractorMixin):
         )
 
     def _validate_configuration(self) -> None:
-        if len(self.component_config.get(CONFIG_FEATURES, [])) % 2 != 1:
+        if len(self.component_config.get(self.CONFIG_FEATURES, [])) % 2 != 1:
             raise ValueError(
                 "Need an odd number of crf feature lists to have a center word."
             )
@@ -275,11 +251,9 @@ class CRFEntityExtractor(GraphComponent, EntityExtractorMixin):
         ]
         dataset = [self._convert_to_crf_tokens(example) for example in entity_examples]
 
-        self.entity_taggers = self.train_model(
-            dataset, self.component_config, self.crf_order
-        )
+        self._train_model(dataset)
 
-        self.persist(dataset)
+        self.persist()
 
         return self._resource
 
@@ -325,9 +299,7 @@ class CRFEntityExtractor(GraphComponent, EntityExtractorMixin):
             if include_tag_features:
                 self._add_tag_to_crf_token(crf_tokens, predictions)
 
-            features = self._crf_tokens_to_features(
-                crf_tokens, self.component_config, include_tag_features
-            )
+            features = self._crf_tokens_to_features(crf_tokens, include_tag_features)
             predictions[tag_name] = entity_tagger.predict_marginals_single(features)
 
         # convert predictions into a list of tags and a list of confidences
@@ -417,25 +389,27 @@ class CRFEntityExtractor(GraphComponent, EntityExtractorMixin):
         **kwargs: Any,
     ) -> CRFEntityExtractor:
         """Loads trained component (see parent class for full docstring)."""
+        import joblib
+
         try:
+            entity_taggers = OrderedDict()
             with model_storage.read_from(resource) as model_dir:
-                dataset = rasa.shared.utils.io.read_json_file(
-                    model_dir / "crf_dataset.json"
-                )
-                crf_order = rasa.shared.utils.io.read_json_file(
-                    model_dir / "crf_order.json"
-                )
-
-                dataset = [
-                    [CRFToken.create_from_dict(token_data) for token_data in sub_list]
-                    for sub_list in dataset
-                ]
+                # We have to load in the same order as we persisted things as otherwise
+                # the predictions might be off
+                file_names = sorted(model_dir.glob("**/*.pkl"))
+                if not file_names:
+                    logger.debug(
+                        "Failed to load model for 'CRFEntityExtractor'. "
+                        "Maybe you did not provide enough training data and "
+                        "no model was trained."
+                    )
+                    return cls(config, model_storage, resource)
 
-                entity_taggers = cls.train_model(dataset, config, crf_order)
+                for file_name in file_names:
+                    name = file_name.stem[1:]
+                    entity_taggers[name] = joblib.load(file_name)
 
-                entity_extractor = cls(config, model_storage, resource, entity_taggers)
-                entity_extractor.crf_order = crf_order
-                return entity_extractor
+                return cls(config, model_storage, resource, entity_taggers)
         except ValueError:
             logger.warning(
                 f"Failed to load {cls.__name__} from model storage. Resource "
@@ -443,29 +417,23 @@ class CRFEntityExtractor(GraphComponent, EntityExtractorMixin):
             )
             return cls(config, model_storage, resource)
 
-    def persist(self, dataset: List[List[CRFToken]]) -> None:
+    def persist(self) -> None:
         """Persist this model into the passed directory."""
-        with self._model_storage.write_to(self._resource) as model_dir:
-            data_to_store = [
-                [token.to_dict() for token in sub_list] for sub_list in dataset
-            ]
+        import joblib
 
-            rasa.shared.utils.io.dump_obj_as_json_to_file(
-                model_dir / "crf_dataset.json", data_to_store
-            )
-            rasa.shared.utils.io.dump_obj_as_json_to_file(
-                model_dir / "crf_order.json", self.crf_order
-            )
+        with self._model_storage.write_to(self._resource) as model_dir:
+            if self.entity_taggers:
+                for idx, (name, entity_tagger) in enumerate(
+                    self.entity_taggers.items()
+                ):
+                    model_file_name = model_dir / f"{idx}{name}.pkl"
+                    joblib.dump(entity_tagger, model_file_name)
 
-    @classmethod
     def _crf_tokens_to_features(
-        cls,
-        crf_tokens: List[CRFToken],
-        config: Dict[str, Any],
-        include_tag_features: bool = False,
+        self, crf_tokens: List[CRFToken], include_tag_features: bool = False
     ) -> List[Dict[Text, Any]]:
         """Convert the list of tokens into discrete features."""
-        configured_features = config[CONFIG_FEATURES]
+        configured_features = self.component_config[self.CONFIG_FEATURES]
         sentence_features = []
 
         for token_idx in range(len(crf_tokens)):
@@ -476,31 +444,28 @@ class CRFEntityExtractor(GraphComponent, EntityExtractorMixin):
             half_window_size = window_size // 2
             window_range = range(-half_window_size, half_window_size + 1)
 
-            token_features = cls._create_features_for_token(
+            token_features = self._create_features_for_token(
                 crf_tokens,
                 token_idx,
                 half_window_size,
                 window_range,
                 include_tag_features,
-                config,
             )
 
             sentence_features.append(token_features)
 
         return sentence_features
 
-    @classmethod
     def _create_features_for_token(
-        cls,
+        self,
         crf_tokens: List[CRFToken],
         token_idx: int,
         half_window_size: int,
         window_range: range,
         include_tag_features: bool,
-        config: Dict[str, Any],
     ) -> Dict[Text, Any]:
         """Convert a token into discrete features including words before and after."""
-        configured_features = config[CONFIG_FEATURES]
+        configured_features = self.component_config[self.CONFIG_FEATURES]
         prefixes = [str(i) for i in window_range]
 
         token_features = {}
@@ -540,13 +505,13 @@ class CRFEntityExtractor(GraphComponent, EntityExtractorMixin):
                         # set in the training data, 'matched' is either 'True' or
                         # 'False' depending on whether the token actually matches the
                         # pattern or not
-                        regex_patterns = cls.function_dict[feature](token)
+                        regex_patterns = self.function_dict[feature](token)
                         for pattern_name, matched in regex_patterns.items():
-                            token_features[
-                                f"{prefix}:{feature}:{pattern_name}"
-                            ] = matched
+                            token_features[f"{prefix}:{feature}:{pattern_name}"] = (
+                                matched
+                            )
                     else:
-                        value = cls.function_dict[feature](token)
+                        value = self.function_dict[feature](token)
                         token_features[f"{prefix}:{feature}"] = value
 
         return token_features
@@ -670,46 +635,38 @@ class CRFEntityExtractor(GraphComponent, EntityExtractorMixin):
 
         return tags
 
-    @classmethod
-    def train_model(
-        cls,
-        df_train: List[List[CRFToken]],
-        config: Dict[str, Any],
-        crf_order: List[str],
-    ) -> OrderedDict[str, CRF]:
+    def _train_model(self, df_train: List[List[CRFToken]]) -> None:
         """Train the crf tagger based on the training data."""
         import sklearn_crfsuite
 
-        entity_taggers = OrderedDict()
+        self.entity_taggers = OrderedDict()
 
-        for tag_name in crf_order:
+        for tag_name in self.crf_order:
             logger.debug(f"Training CRF for '{tag_name}'.")
 
             # add entity tag features for second level CRFs
             include_tag_features = tag_name != ENTITY_ATTRIBUTE_TYPE
             X_train = (
-                cls._crf_tokens_to_features(sentence, config, include_tag_features)
+                self._crf_tokens_to_features(sentence, include_tag_features)
                 for sentence in df_train
             )
             y_train = (
-                cls._crf_tokens_to_tags(sentence, tag_name) for sentence in df_train
+                self._crf_tokens_to_tags(sentence, tag_name) for sentence in df_train
             )
 
             entity_tagger = sklearn_crfsuite.CRF(
                 algorithm="lbfgs",
                 # coefficient for L1 penalty
-                c1=config["L1_c"],
+                c1=self.component_config["L1_c"],
                 # coefficient for L2 penalty
-                c2=config["L2_c"],
+                c2=self.component_config["L2_c"],
                 # stop earlier
-                max_iterations=config["max_iterations"],
+                max_iterations=self.component_config["max_iterations"],
                 # include transitions that are possible, but not observed
                 all_possible_transitions=True,
             )
             entity_tagger.fit(X_train, y_train)
 
-            entity_taggers[tag_name] = entity_tagger
+            self.entity_taggers[tag_name] = entity_tagger
 
             logger.debug("Training finished.")
-
-        return entity_taggers