Add Django management command to train chat bot

chatterbot/chatterbot.py

@@ -2,7 +2,6 @@
 from .adapters.logic import LogicAdapter, MultiLogicAdapter
 from .adapters.input import InputAdapter
 from .adapters.output import OutputAdapter
-from .conversation import Statement, Response
 from .utils.queues import ResponseQueue
 from .utils.module_loading import import_module
 import logging
@@ -67,6 +66,7 @@ def __init__(self, name, **kwargs):
         trainer = kwargs.get('trainer', 'chatterbot.trainers.Trainer')
         TrainerClass = import_module(trainer)
         self.trainer = TrainerClass(self.storage, **kwargs)
+        self.training_data = kwargs.get('training_data')
 
         self.logger = kwargs.get('logger', logging.getLogger(__name__))
 
@@ -228,6 +228,8 @@ def learn_response(self, statement):
         """
         Learn that the statement provided is a valid response.
         """
+        from .conversation import Response
+
         previous_statement = self.get_last_response_statement()
 
         if previous_statement:

chatterbot/ext/django_chatterbot/management/commands/train.py

@@ -0,0 +1,23 @@
+from django.core.management.base import BaseCommand
+
+
+class Command(BaseCommand):
+    help = 'Trains the database used by the chat bot'
+    can_import_settings = True
+
+    def handle(self, *args, **options):
+        from chatterbot import ChatBot
+        from chatterbot.ext.django_chatterbot import settings
+
+        chatterbot = ChatBot(**settings.CHATTERBOT)
+
+        chatterbot.train(chatterbot.training_data)
+
+        # Django 1.8 does not define SUCCESS
+        if hasattr(self.style, 'SUCCESS'):
+            style = self.style.SUCCESS
+        else:
+            style = self.style.NOTICE
+
+        training_class = chatterbot.trainer.__class__.__name__
+        self.stdout.write(style('ChatterBot trained using "%s"' % training_class))

chatterbot/trainers.py

@@ -1,13 +1,11 @@
 from .conversation import Statement, Response
-from .corpus import Corpus
 import logging
 
 
 class Trainer(object):
 
     def __init__(self, storage, **kwargs):
         self.storage = storage
-        self.corpus = Corpus()
         self.logger = logging.getLogger(__name__)
 
     def train(self, *args, **kwargs):
@@ -76,9 +74,20 @@ def train(self, conversation):
 
 class ChatterBotCorpusTrainer(Trainer):
 
+    def __init__(self, storage, **kwargs):
+        super(ChatterBotCorpusTrainer, self).__init__(storage, **kwargs)
+        from .corpus import Corpus
+
+        self.corpus = Corpus()
+
     def train(self, *corpora):
         trainer = ListTrainer(self.storage)
 
+        # Allow a list of coupora to be passed instead of arguments
+        if len(corpora) == 1:
+            if isinstance(corpora[0], list):
+                corpora = corpora[0]
+
         for corpus in corpora:
             corpus_data = self.corpus.load_corpus(corpus)
             for data in corpus_data:

docs/django/index.rst

@@ -10,6 +10,7 @@ Django applications.
    :maxdepth: 2
 
    settings
+   training
    views
 
 Installation

docs/django/settings.rst

@@ -12,6 +12,10 @@ You can edit the ChatterBot configuration through your Django settings.py file.
            'chatterbot.adapters.logic.MathematicalEvaluation',
            'chatterbot.adapters.logic.TimeLogicAdapter',
            'chatterbot.adapters.logic.ClosestMatchAdapter'
+       ],
+       'trainer': 'chatterbot.trainers.ChatterBotCorpusTrainer',
+       'training_data': [
+            'chatterbot.corpus.english.greetings'
        ]
    }
 

docs/django/training.rst

@@ -0,0 +1,30 @@
+========
+Training
+========
+
+Management command
+==================
+
+When using ChatterBot with Django, the training process can be
+executed by running the training management command. 
+
+.. code-block:: bash
+
+   python manage.py train
+
+Training settings
+=================
+
+You can specify any data that you want to be passed to the chat bot
+trainer in the :code:`training_data` parameter in your :code:`CHATTERBOT`
+Django settings.
+
+.. code-block:: python
+
+   CHATTERBOT = {
+       # ...
+       'trainer': 'chatterbot.trainers.ChatterBotCorpusTrainer',
+       'training_data': [
+            'chatterbot.corpus.english.greetings'
+       ]
+   }

examples/django_app/example_app/settings.py

@@ -33,7 +33,11 @@
 # ChatterBot settings
 
 CHATTERBOT = {
-    'name': 'Django ChatterBot Example'
+    'name': 'Django ChatterBot Example',
+    'trainer': 'chatterbot.trainers.ChatterBotCorpusTrainer',
+    'training_data': [
+        'chatterbot.corpus.english.greetings'
+    ]
 }
 
 MIDDLEWARE_CLASSES = (

examples/django_app/tests/test_commands.py

@@ -0,0 +1,21 @@
+from django.core.management import call_command
+from django.test import TestCase
+from django.utils.six import StringIO
+from chatterbot.ext.django_chatterbot.models import Statement
+
+
+class TrainCommandTestCase(TestCase):
+
+    def test_command_output(self):
+        out = StringIO()
+        call_command('train', stdout=out)
+        self.assertIn('ChatterBot trained', out.getvalue())
+
+    def test_command_data_argument(self):
+        out = StringIO()
+        statements_before = Statement.objects.exists()
+        call_command('train', stdout=out)
+        statements_after = Statement.objects.exists()
+
+        self.assertFalse(statements_before)
+        self.assertTrue(statements_after)

setup.py

@@ -33,6 +33,8 @@
         'chatterbot.ext',
         'chatterbot.ext.django_chatterbot',
         'chatterbot.ext.django_chatterbot.migrations',
+        'chatterbot.ext.django_chatterbot.management',
+        'chatterbot.ext.django_chatterbot.management.commands',
         'chatterbot.utils'
     ],
     package_dir={'chatterbot': 'chatterbot'},