markov.py 2.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182
  1. import os.path
  2. import re
  3. import atexit
  4. import ujson
  5. import markovify
  6. from config import config
  7. class Markov:
  8. def __init__(self):
  9. self.counter = 0
  10. self.corpus = []
  11. self.chain = None
  12. self.load()
  13. atexit.register(self.save)
  14. @property
  15. def is_ready(self):
  16. return self.chain is not None
  17. def generate(self):
  18. words = self.chain.walk()
  19. if not words:
  20. return self.generate()
  21. text = " ".join(words)
  22. text = re.sub(r"(?:^| )?((\?\.\.)|(\.{2,})|(\!{2,})|(\?{2,})|([.?!,:;\(\)\"'\$\+\-–—…]))(?: |$)", r"\1 ", text)
  23. text = text.strip()
  24. return text
  25. def rebuild(self):
  26. self.chain = markovify.Chain(self.corpus, config.MARKOV_STATE_SIZE).compile()
  27. self.counter = 0
  28. def extend_corpus(self, text):
  29. text = text.strip()
  30. if not text:
  31. return
  32. text = text.replace("\n", " ")
  33. text = re.sub(r"(@[a-z0-9_]+,?)", "", text)
  34. text = re.sub("https?:\\/\\/(?:www\\.)?[-a-zA-Z0-9@:%._\\+~#=]{1,256}\\.[a-zA-Z0-9()]{1,6}\\b(?:[-a-zA-Z0-9()@:%_\\+.~#?&\\/=]*)", "", text)
  35. text = re.sub(r"((\?\.\.)|(\.{2,})|(\!{2,})|(\?{2,})|[.?!,:;\(\)\"'\$\+\-–—…])", r" \1 ", text)
  36. text = text.split(" ")
  37. text = map(lambda word: word.strip(), text)
  38. text = filter(bool, text)
  39. text = list(text)
  40. if text not in self.corpus:
  41. self.corpus.insert(0, text)
  42. if len(self.corpus) > config.MARKOV_CORPUS_SIZE:
  43. self.corpus = self.corpus[: config.MARKOV_CORPUS_SIZE]
  44. self.counter += 1
  45. if self.counter % config.MARKOV_REBUILD_RATE == 0:
  46. self.rebuild()
  47. def load(self):
  48. if os.path.isfile(config.MARKOV_CHAIN_PATH):
  49. with open(config.MARKOV_CHAIN_PATH, "r") as f:
  50. self.chain = markovify.Chain.from_json(f.read())
  51. if os.path.isfile(config.MARKOV_CORPUS_PATH):
  52. with open(config.MARKOV_CORPUS_PATH, "r") as f:
  53. self.corpus = ujson.load(f)
  54. def save(self):
  55. if self.chain:
  56. with open(config.MARKOV_CHAIN_PATH, "w") as f:
  57. f.write(self.chain.to_json())
  58. with open(config.MARKOV_CORPUS_PATH, "w") as f:
  59. ujson.dump(self.corpus, f)