Advertisement
Guest User

Untitled

a guest
Jul 23rd, 2016
75
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 8.46 KB | None | 0 0
  1. # Copyright (C) 2014 Peter Teichman
  2.  
  3. import atexit
  4. import logging
  5. import os
  6. import re
  7. import readline
  8. import Stemmer
  9. import sys
  10. import time
  11.  
  12. from .bot import Runner
  13. from .brain import Brain
  14.  
  15. log = logging.getLogger("cobe")
  16.  
  17.  
  18. class InitCommand:
  19. @classmethod
  20. def add_subparser(cls, parser):
  21. subparser = parser.add_parser("init", help="Initialize a new brain")
  22.  
  23. subparser.add_argument("--force", action="store_true")
  24. subparser.add_argument("--order", type=int, default=3)
  25. subparser.add_argument("--megahal", action="store_true",
  26. help="Use MegaHAL-compatible tokenizer")
  27. subparser.set_defaults(run=cls.run)
  28.  
  29. @staticmethod
  30. def run(args):
  31. filename = args.brain
  32.  
  33. if os.path.exists(filename):
  34. if args.force:
  35. os.remove(filename)
  36. else:
  37. log.error("%s already exists!", filename)
  38. return
  39.  
  40. tokenizer = None
  41. if args.megahal:
  42. tokenizer = "MegaHAL"
  43.  
  44. Brain.init(filename, order=args.order, tokenizer=tokenizer)
  45.  
  46.  
  47. def progress_generator(filename):
  48. s = os.stat(filename)
  49. size_left = s.st_size
  50.  
  51. fd = open(filename)
  52. for line in fd.xreadlines():
  53. size_left = size_left - len(line)
  54. progress = 100 * (1. - (float(size_left) / float(s.st_size)))
  55.  
  56. yield line, progress
  57.  
  58. fd.close()
  59.  
  60.  
  61. class LearnCommand:
  62. @classmethod
  63. def add_subparser(cls, parser):
  64. subparser = parser.add_parser("learn", help="Learn a file of text")
  65. subparser.add_argument("file", nargs="+")
  66. subparser.set_defaults(run=cls.run)
  67.  
  68. @staticmethod
  69. def run(args):
  70. b = Brain(args.brain)
  71. b.start_batch_learning()
  72.  
  73. for filename in args.file:
  74. now = time.time()
  75. print filename
  76.  
  77. count = 0
  78. for line, progress in progress_generator(filename):
  79. show_progress = ((count % 1000) == 0)
  80.  
  81. if show_progress:
  82. elapsed = time.time() - now
  83. sys.stdout.write("\r%.0f%% (%d/s)" % (progress,
  84. count / elapsed))
  85. sys.stdout.flush()
  86.  
  87. b.learn(line.strip())
  88. count = count + 1
  89.  
  90. if (count % 10000) == 0:
  91. b.graph.commit()
  92.  
  93. elapsed = time.time() - now
  94. print "\r100%% (%d/s)" % (count / elapsed)
  95.  
  96. b.stop_batch_learning()
  97.  
  98.  
  99. class LearnIrcLogCommand:
  100. @classmethod
  101. def add_subparser(cls, parser):
  102. subparser = parser.add_parser("learn-irc-log",
  103. help="Learn a file of IRC log text")
  104. subparser.add_argument("-i", "--ignore-nick", action="append",
  105. dest="ignored_nicks",
  106. help="Ignore an IRC nick")
  107. subparser.add_argument("-o", "--only-nick", action="append",
  108. dest="only_nicks",
  109. help="Only learn from specified nicks")
  110. subparser.add_argument("-r", "--reply-to", action="append",
  111. help="Reply (invisibly) to things said "
  112. "to specified nick")
  113. subparser.add_argument("file", nargs="+")
  114. subparser.set_defaults(run=cls.run)
  115.  
  116. @classmethod
  117. def run(cls, args):
  118. b = Brain(args.brain)
  119. b.start_batch_learning()
  120.  
  121. for filename in args.file:
  122. now = time.time()
  123. print filename
  124.  
  125. count = 0
  126. for line, progress in progress_generator(filename):
  127. show_progress = ((count % 100) == 0)
  128.  
  129. if show_progress:
  130. elapsed = time.time() - now
  131. sys.stdout.write("\r%.0f%% (%d/s)" % (progress,
  132. count / elapsed))
  133. sys.stdout.flush()
  134.  
  135. count = count + 1
  136.  
  137. if (count % 1000) == 0:
  138. b.graph.commit()
  139.  
  140. parsed = cls._parse_irc_message(line.strip(),
  141. args.ignored_nicks,
  142. args.only_nicks)
  143. if parsed is None:
  144. continue
  145.  
  146. to, msg = parsed
  147. b.learn(msg)
  148.  
  149. if args.reply_to is not None and to in args.reply_to:
  150. b.reply(msg)
  151.  
  152. elapsed = time.time() - now
  153. print "\r100%% (%d/s)" % (count / elapsed)
  154.  
  155. b.stop_batch_learning()
  156.  
  157. @staticmethod
  158. def _parse_irc_message(msg, ignored_nicks=None, only_nicks=None):
  159. # only match lines of the form "HH:MM <nick> message"
  160. match = re.match("\d+:\d+\s+<(.+?)>\s+(.*)", msg)
  161. if not match:
  162. return None
  163.  
  164. nick = match.group(1)
  165. msg = match.group(2)
  166.  
  167. if ignored_nicks is not None and nick in ignored_nicks:
  168. return None
  169.  
  170. if only_nicks is not None and nick not in only_nicks:
  171. return None
  172.  
  173. to = None
  174.  
  175. # strip "username: " at the beginning of messages
  176. match = re.search("^(\S+)[,:]\s+(\S.*)", msg)
  177. if match:
  178. to = match.group(1)
  179. msg = match.group(2)
  180.  
  181. # strip kibot style '"asdf" --user, 06-oct-09' quotes
  182. msg = re.sub("\"(.*)\" --\S+,\s+\d+-\S+-\d+",
  183. lambda m: m.group(1), msg)
  184.  
  185. return to, msg
  186.  
  187.  
  188. class ConsoleCommand:
  189. @classmethod
  190. def add_subparser(cls, parser):
  191. subparser = parser.add_parser("console", help="Interactive console")
  192. subparser.set_defaults(run=cls.run)
  193.  
  194. @staticmethod
  195. def run(args):
  196. b = Brain(args.brain)
  197.  
  198. history = os.path.expanduser("~/.cobe_history")
  199. try:
  200. readline.read_history_file(history)
  201. except IOError:
  202. pass
  203. atexit.register(readline.write_history_file, history)
  204.  
  205. while True:
  206. try:
  207. cmd = raw_input("> ")
  208. except EOFError:
  209. print
  210. sys.exit(0)
  211.  
  212. b.learn(cmd)
  213. print b.reply(cmd).encode("utf-8")
  214.  
  215.  
  216. class IrcClientCommand:
  217. @classmethod
  218. def add_subparser(cls, parser):
  219. subparser = parser.add_parser("irc-client",
  220. help="IRC client [requires twisted]")
  221. subparser.add_argument("-s", "--server", required=True,
  222. help="IRC server hostname")
  223. subparser.add_argument("-p", "--port", type=int, default=6667,
  224. help="IRC server port")
  225. subparser.add_argument("-n", "--nick", default="cobe",
  226. help="IRC nick")
  227. subparser.add_argument("-c", "--channel", required=True,
  228. help="IRC channel")
  229. subparser.add_argument("-l", "--log-channel",
  230. help="IRC channel for logging")
  231. subparser.add_argument("-i", "--ignore-nick", action="append",
  232. dest="ignored_nicks",
  233. help="Ignore an IRC nick")
  234. subparser.add_argument("-o", "--only-nick", action="append",
  235. dest="only_nicks",
  236. help="Only learn from a specific IRC nick")
  237.  
  238. subparser.set_defaults(run=cls.run)
  239.  
  240. @staticmethod
  241. def run(args):
  242. b = Brain(args.brain)
  243.  
  244. Runner().run(b, args)
  245.  
  246.  
  247. class SetStemmerCommand:
  248. @classmethod
  249. def add_subparser(cls, parser):
  250. subparser = parser.add_parser("set-stemmer",
  251. help="Configure a stemmer")
  252.  
  253. subparser.set_defaults(run=cls.run)
  254.  
  255. subparser.add_argument("language", choices=Stemmer.algorithms(),
  256. help="Stemmer language")
  257.  
  258. @staticmethod
  259. def run(args):
  260. b = Brain(args.brain)
  261.  
  262. b.set_stemmer(args.language)
  263.  
  264.  
  265. class DelStemmerCommand:
  266. @classmethod
  267. def add_subparser(cls, parser):
  268. subparser = parser.add_parser("del-stemmer", help="Delete the stemmer")
  269. subparser.set_defaults(run=cls.run)
  270.  
  271. @staticmethod
  272. def run(args):
  273. b = Brain(args.brain)
  274.  
  275. b.del_stemmer()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement