modules/markov/markov.py

   1 #markov module for scrappy
   2
   3 #the inspiration for this code comes from
   4 #"The Practice of Programming", so some of the variable
   5 #names and such will be similar.
   6
   7 import random
   8 import pickle
   9 import re
  10
  11 nickmatch = None
  12 statetab = {}
  13
  14 w1 = w2 = "\n"
  15
  16 def init(scrap):
  17     global nickmatch
  18
  19     scrap.register_event("msg", markov_learn)
  20     scrap.register_event("msg", markov_talk)
  21     scrap.register_event("msg", markov_load)
  22     scrap.register_event("msg", markov_dump)
  23     scrap.register_event("msg", markov_stats)
  24
  25     nickmatch = re.compile(scrap.nickname)
  26
  27     random.seed()
  28
  29 def markov_stats(c,list,bot):
  30     global statetab
  31
  32     cmd = list[4].split(" ")[0]
  33
  34     if cmd == "markov_stats":
  35         c.privmsg(list[5], "words: %d" % len(statetab[("\n","\n")]))
  36         c.privmsg(list[5], "chains: %d" % len(statetab.items()))
  37
  38 #loads in a previously pickled saved state
  39 def markov_load(c,list,bot):
  40     global statetab
  41
  42     cmd = list[4].split(" ")[0]
  43
  44     if cmd == "mkload":
  45         fp = list[4].split(" ")[1]
  46
  47         try:
  48                 pkfile = open(fp,"r")
  49                 statetab = pickle.load(pkfile)
  50         except IOError:
  51             print "Could not load db: Doesn't exist\n"
  52
  53
  54 #pickles out the state to a file
  55 def markov_dump(c,list,bot):
  56     global statetab
  57
  58     cmd = list[4].split(" ")[0]
  59
  60     if cmd == "mkdump":
  61         fp = list[4].split(" ")[1]
  62
  63         pkfile = open(fp,"w+")
  64
  65         pickle.dump(statetab,pkfile)
  66
  67 def markov_learn(c,list,bot):
  68     """ Should not be called directly """
  69
  70     words = [x for x in list[4].split(" ") if not x.isspace()]
  71
  72     global statetab
  73     global w1
  74     global w2
  75
  76     w1 = w2 = "\n"
  77
  78     #go through every word and put them in a hash table.
  79     #EX the sentence "Mary had a little lamb"
  80     #first iteration, w1 and w2 are both empty.
  81     #statetab[w1][w2] doesn't exist, so make it and set
  82     #statetab[""][""] to Mary.
  83     #
  84     #Then, set w1 to w2 and w2 to i, so the chain moves forward.
  85     for i in words:
  86         statetab.setdefault((w1,w2),[]).append(i)
  87         w1,w2 = w2, i
  88
  89     statetab.setdefault((w1,w2),[]).append("\n")
  90
  91     if nickmatch.search(list[4]) and bot.autorep == 1 and random.randint(0,3) == 0:
  92         tmp = emit_chain(random.choice(list[4].split(" ")))
  93
  94         if len(tmp) <= 2:
  95             return
  96
  97         c.privmsg(list[5], "%s: %s" % (list[0],tmp))
  98         return
  99
 100     #randomly reply
 101     if random.randint(0,15) == 0 and bot.talk == 1:
 102         c.privmsg(list[5], "%s" % (emit_chain(random.choice(list[4].split(" ")))))
 103
 104 def emit_chain(key):
 105     global statetab
 106     global w1
 107     global w2
 108
 109     i = 0
 110
 111     w1 = w2 = "\n"
 112
 113     newword = ""
 114
 115     #make the first word the key if its not a space
 116 #    if(key != " "):
 117 #        retval = key + " "
 118 #    else:
 119     retval = ""
 120
 121     if key != " ":
 122         w2 = key
 123
 124     while 1:
 125         try:
 126             newword = random.choice(statetab[(w1,w2)])
 127         except KeyError:
 128             return retval
 129
 130         retval = retval + newword + " "
 131         w1,w2 = w2,newword
 132
 133         i = i + 1
 134
 135         #max of rand words if we don't hit a space or other error
 136         if i >= random.randint(5,50):
 137             return retval
 138
 139     return retval
 140
 141 def markov_talk(c,list,bot):
 142     """ Makes the markov chain talk to you """
 143
 144     cmd = list[4].split(" ")[0]
 145
 146     try:
 147         key = list[4].split(" ")[1]
 148     except IndexError:
 149         key = " "
 150
 151     if list[3] and cmd == "talk":
 152         tmp = emit_chain(key);
 153         if len(tmp) <= 2:
 154             return
 155         c.privmsg(list[5],"%s" % tmp)