From 6e2141679140d52754d766aed0dce2664492607f Mon Sep 17 00:00:00 2001 From: "Brian S. Stephan" Date: Wed, 16 Nov 2016 12:10:42 -0600 Subject: [PATCH] markov: import generic text files --- dr_botzo/markov/admin.py | 67 +++++++++++++------ dr_botzo/markov/forms.py | 17 +++-- .../migrations/0003_auto_20161112_2348.py | 18 +++++ dr_botzo/markov/models.py | 2 +- 4 files changed, 79 insertions(+), 25 deletions(-) create mode 100644 dr_botzo/markov/migrations/0003_auto_20161112_2348.py diff --git a/dr_botzo/markov/admin.py b/dr_botzo/markov/admin.py index aef6201..becee38 100644 --- a/dr_botzo/markov/admin.py +++ b/dr_botzo/markov/admin.py @@ -19,7 +19,7 @@ admin.site.register(MarkovTarget) admin.site.register(MarkovState) -@permission_required('import_log_file', raise_exception=True) +@permission_required('import_text_file', raise_exception=True) def import_file(request): """Accept a file upload and turn it into markov stuff. @@ -30,31 +30,58 @@ def import_file(request): if request.method == 'POST': form = LogUploadForm(request.POST, request.FILES) if form.is_valid(): - log_file = request.FILES['log_file'] - context = form.cleaned_data['context'] - ignores = form.cleaned_data['ignore_nicks'].split(',') - strips = form.cleaned_data['strip_prefixes'].split(' ') + if form.cleaned_data['text_file_format'] == LogUploadForm.FILE_FORMAT_WEECHAT: + text_file = request.FILES['text_file'] + context = form.cleaned_data['context'] + ignores = form.cleaned_data['ignore_nicks'].split(',') + strips = form.cleaned_data['strip_prefixes'].split(' ') - whos = [] - for line in log_file: - log.debug(line) - (timestamp, who, what) = line.decode('utf-8').split('\t', 2) + whos = [] + for line in text_file: + log.debug(line) + (timestamp, who, what) = line.decode('utf-8').split('\t', 2) - if who in ('-->', '<--', '--', ' *'): - continue + if who in ('-->', '<--', '--', ' *'): + continue - if who in ignores: - continue + if who in ignores: + continue - whos.append(who) + whos.append(who) - # this is a line we probably care about now - what = [x for x in what.rstrip().split(' ') if x not in strips] - markovlib.learn_line(' '.join(what), context) - log.debug("learned") + # this is a line we probably care about now + what = [x for x in what.rstrip().split(' ') if x not in strips] + markovlib.learn_line(' '.join(what), context) + log.debug("learned") - log.debug(set(whos)) - form = LogUploadForm() + log.debug(set(whos)) + form = LogUploadForm() + elif form.cleaned_data['text_file_format'] == LogUploadForm.FILE_FORMAT_RAW_TEXT: + text_file = request.FILES['text_file'] + context = form.cleaned_data['context'] + + k1 = MarkovState._start1 + k2 = MarkovState._start2 + for line in text_file: + for word in [x for x in line.decode('utf-8') .rstrip().split(' ')]: + log.info(word) + if word: + state, created = MarkovState.objects.get_or_create(context=context, k1=k1, + k2=k2, v=word) + state.count += 1 + state.save() + + if word[-1] in ['.', '?', '!']: + state, created = MarkovState.objects.get_or_create(context=context, k1=k2, + k2=word, v=MarkovState._stop) + state.count += 1 + state.save() + + k1 = MarkovState._start1 + k2 = MarkovState._start2 + else: + k1 = k2 + k2 = word else: form = LogUploadForm() diff --git a/dr_botzo/markov/forms.py b/dr_botzo/markov/forms.py index 4b442e5..a07093a 100644 --- a/dr_botzo/markov/forms.py +++ b/dr_botzo/markov/forms.py @@ -2,7 +2,7 @@ import logging -from django.forms import Form, CharField, FileField, ModelChoiceField +from django.forms import Form, CharField, ChoiceField, FileField, ModelChoiceField from markov.models import MarkovContext @@ -13,11 +13,20 @@ class LogUploadForm(Form): """Accept a file upload that will be imported into Markov stuff.""" - log_file = FileField(help_text="Weechat log format.") + FILE_FORMAT_WEECHAT = 'WEECHAT' + FILE_FORMAT_RAW_TEXT = 'RAW' + + FILE_FORMAT_CHOICES = ( + (FILE_FORMAT_WEECHAT, "Weechat"), + (FILE_FORMAT_RAW_TEXT, "Raw text file"), + ) + + text_file = FileField() + text_file_format = ChoiceField(choices=FILE_FORMAT_CHOICES) context = ModelChoiceField(queryset=MarkovContext.objects.all()) - ignore_nicks = CharField(help_text="Comma-separated list of nicks to ignore.", + ignore_nicks = CharField(help_text="Comma-separated list of nicks to ignore. For Weechat logs.", required=False) - strip_prefixes = CharField(help_text="Space-separated list of line prefixes to strip.", + strip_prefixes = CharField(help_text="Space-separated list of line prefixes to strip. For Weechat logs.", required=False) diff --git a/dr_botzo/markov/migrations/0003_auto_20161112_2348.py b/dr_botzo/markov/migrations/0003_auto_20161112_2348.py new file mode 100644 index 0000000..bededdd --- /dev/null +++ b/dr_botzo/markov/migrations/0003_auto_20161112_2348.py @@ -0,0 +1,18 @@ +# -*- coding: utf-8 -*- +from __future__ import unicode_literals + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('markov', '0002_auto_20150514_2317'), + ] + + operations = [ + migrations.AlterModelOptions( + name='markovstate', + options={'permissions': set([('import_text_file', 'Can import states from a text file'), ('teach_line', 'Can teach lines')])}, + ), + ] diff --git a/dr_botzo/markov/models.py b/dr_botzo/markov/models.py index 560615f..c78a19b 100644 --- a/dr_botzo/markov/models.py +++ b/dr_botzo/markov/models.py @@ -59,7 +59,7 @@ class MarkovState(models.Model): ['context', 'v'], ] permissions = { - ('import_log_file', "Can import states from a log file"), + ('import_text_file', "Can import states from a text file"), ('teach_line', "Can teach lines"), } unique_together = ('context', 'k1', 'k2', 'v')