modules/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
  10 statetab = {}
  11
  12 w1 = w2 = "\n"
  13
  14 def __init__(scrap):
  15     scrap.register_msg_event(markov_learn)
  16     scrap.register_msg_event(markov_talk)
  17     scrap.register_msg_event(markov_load)
  18     scrap.register_msg_event(markov_dump)
  19
  20     random.seed()
  21
  22 #loads in a previously pickled saved state
  23 def markov_load(c,list,bot):
  24     global statetab
  25
  26     cmd = list[4].split(" ")[0]
  27
  28     if cmd == "mkload":
  29         fp = list[4].split(" ")[1]
  30
  31         pkfile = open(fp,"r")
  32
  33         statetab = pickle.load(pkfile)
  34
  35 #pickles out the state to a file
  36 def markov_dump(c,list,bot):
  37     global statetab
  38
  39     cmd = list[4].split(" ")[0]
  40
  41     if cmd == "mkdump":
  42         fp = list[4].split(" ")[1]
  43
  44         pkfile = open(fp,"w+")
  45
  46         pickle.dump(statetab,pkfile)
  47
  48 def markov_learn(c,list,bot):
  49     """ Should not be called directly """
  50
  51     words = [x for x in list[4].split(" ") if not x.isspace()]
  52
  53     global statetab
  54     global w1
  55     global w2
  56
  57     w1 = w2 = "\n"
  58
  59     #go through every word and put them in a hash table.
  60     #EX the sentence "Mary had a little lamb"
  61     #first iteration, w1 and w2 are both empty.
  62     #statetab[w1][w2] doesn't exist, so make it and set
  63     #statetab[""][""] to Mary.
  64     #
  65     #Then, set w1 to w2 and w2 to i, so the chain moves forward.
  66     for i in words:
  67         statetab.setdefault((w1,w2),[]).append(i)
  68         w1,w2 = w2, i
  69
  70     statetab.setdefault((w1,w2),[]).append("\n")
  71
  72     #randomly reply
  73     if random.randint(0,20) == 0 and bot.talk == 1:
  74         c.privmsg(list[5], "%s: %s" % (list[0],
  75             emit_chain(random.choice(list[4].split(" ")))))
  76
  77 def emit_chain(key):
  78     global statetab
  79     global w1
  80     global w2
  81
  82     i = 0
  83
  84     w1 = w2 = "\n"
  85
  86     newword = ""
  87
  88     #make the first word the key if its not a space
  89     if(key != " "):
  90         retval = key + " "
  91     else:
  92         retval = ""
  93
  94     if key != " ":
  95         w2 = key
  96
  97     while 1:
  98         if retval.isspace():
  99             emit_chain(key)
 100         try:
 101             newword = random.choice(statetab[(w1,w2)])
 102         except KeyError:
 103             return retval;
 104
 105         retval = retval + newword + " "
 106         w1,w2 = w2,newword
 107
 108         i = i + 1
 109
 110         #max of rand words if we don't hit a space or other error
 111         if i >= random.randint(5,50):
 112             return retval
 113
 114     return retval
 115
 116 def markov_talk(c,list,bot):
 117     """ Makes the markov chain talk to you """
 118
 119     cmd = list[4].split(" ")[0]
 120
 121     try:
 122         key = list[4].split(" ")[1]
 123     except IndexError:
 124         key = " "
 125
 126     if list[3] and cmd == "talk":
 127         c.privmsg(list[5],"%s" % emit_chain(key))