summaryrefslogtreecommitdiff
path: root/chain.py
blob: 895ce1c66f8fa47a17953e4a95c4474ef91103fd (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
import markovify
import re
import spacy

SEPARATOR = r"\s*@@note@@\s*"

nlp = spacy.load("en_core_web_sm")

class Text(markovify.Text):
    def word_split(self, sentence):
        ret = []

        for word in nlp(sentence):
            if word.pos_ == 'PUNCT':
                continue

            ret.append("::".join((word.orth_, word.pos_)))

        return ret

    def word_join(self, words):
        sentence = " ".join(word.split("::")[0] for word in words)
        return sentence
    
    def sentence_split(self, text):
        return re.split(SEPARATOR, text)