Allow user the capability to choose the classifier used by an Agent.

loristns · Jun 22, 2018 · da87d88 · da87d88
1 parent e0c4e9f
commit da87d88
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/kadot/bot_engine.py b/kadot/bot_engine.py
@@ -66,6 +66,7 @@ class Agent(SavedObject):
 
     def __init__(self,
                  word_vectors: Optional[VectorDict] = None,
+                 classifier=NeuralClassifier,
                  tokenizer: Callable[..., Tokens] = regex_tokenizer
                  ):
         """
@@ -75,6 +76,7 @@ def __init__(self,
         :param tokenizer: the word tokenizer to use.
         """
 
+        self.classifier_fn = classifier
         self.classifier = None
         self.tokenizer = tokenizer
         self.word_vectors = word_vectors
@@ -168,7 +170,7 @@ def _get_training_dataset(self):
         return training_dataset
 
     def train(self):
-        self.classifier = NeuralClassifier(
+        self.classifier = self.classifier_fn(
             self._get_training_dataset(),
             word_vectors=self.word_vectors
         )