git.oblomov.eu Git - rbot-mark/blob - mark2.rb

   1 #! /usr/bin/ruby -w
   2 # vim: set sw=2 et:
   3 # Author: Giuseppe Bilotta <giuseppe.bilotta@gmail.com>
   4 # New markov chain plugin
   5
   6 class Array
   7   def butlast
   8     self[0,self.size-1]
   9   end
  10   def butfirst
  11     self[1,self.size]
  12   end
  13 end
  14
  15 class ChanceHash
  16
  17   def initialize
  18     @hash = Hash.new(0)
  19     @picker = {}
  20     @total = 0
  21     @valid_pick = false
  22   end
  23
  24   def increase(el)
  25     if @hash.key?(el)
  26       @hash[el] += 1
  27     else
  28       @hash[el] = 1
  29     end
  30     @valid_pick = false
  31     return @hash[el]
  32   end
  33
  34   def decrease(el)
  35     if @hash.key?(el)
  36       @hash[el] -= 1
  37       @hash.delete(el) if @hash[el] == 0
  38     end
  39     @valid_pick = false
  40     return @hash[el]
  41   end
  42
  43   def make_picker
  44     @picker.clear
  45     total = 0
  46     @hash.each { |el, ch|
  47       total += ch
  48       @picker[total] = el
  49     }
  50     @total = total
  51     @valid_pick = true
  52   end
  53
  54   def random
  55     make_picker unless @valid_pick
  56     pick = rand(@total)
  57     @picker.each { |ch, el|
  58       return el if pick < ch
  59     }
  60   end
  61 end
  62
  63 class MarkovChainer
  64   # Maximum depth
  65   MAX_ORDER = 5
  66
  67   # Word or nonword regexp:
  68   # can be used to scan a string splitting it into
  69   # words and nonwords.
  70   WNW = /\w+|\W/u
  71
  72   def initialize
  73     # mkv[i] holds the chains of order i
  74     @mkv = Array.new
  75
  76     MAX_ORDER.times { |i|
  77       @mkv[i] = {}
  78     }
  79
  80     # Each chain is in the form
  81     # [:array, :of, :symbols] => {
  82     #   :word => [chance before, chance after]
  83     #   :word => [chance before, chance after]
  84     # }
  85     # except for order 0, which is just a hash of
  86     # {:word => chance}
  87   end
  88
  89   def add_one(sym)
  90     s = sym.to_sym rescue nil
  91     if @mkv[0].has_key?(s)
  92       @mkv[0][s] += 1
  93     else
  94       @mkv[0][s] = 1
  95     end
  96   end
  97
  98   def add_before(array, prev)
  99     raise "Not enough words in new data" if array.empty?
 100     raise "Too many words in new data" if array.size > MAX_ORDER
 101     size = array.size
 102     if @mkv[size].has_key?(array)
 103       h = @mkv[size][array]
 104       if h.has_key?(prev)
 105        h[prev][0] += 1
 106       else
 107        h[prev] = [1,0]
 108       end
 109     else
 110       @mkv[size][array.dup] = { prev => [1, 0] }
 111     end
 112   end
 113
 114   def add_after(array, nxt)
 115     raise "Not enough words in new data" if array.empty?
 116     raise "Too many words in new data" if array.size > MAX_ORDER
 117     size = array.size
 118     if @mkv[size].has_key?(array)
 119       h = @mkv[size][array]
 120       if h.has_key?(nxt)
 121        h[nxt][1] += 1
 122       else
 123        h[nxt] = [0,1]
 124       end
 125     else
 126       @mkv[size][array.dup] = { nxt => [0, 1] }
 127     end
 128   end
 129
 130   def add_multi(array)
 131     raise "Too many words in new data" if array.size > MAX_ORDER + 1
 132     add_before(array.butfirst, array.first)
 133     add_after(array.butlast, array.last)
 134   end
 135
 136   def add(*data)
 137     if data.size == 1
 138       add_one(data.first)
 139     else
 140       add_multi(data)
 141     end
 142   end
 143
 144   def simple_learn(text)
 145     syms = text.scan(WNW).map { |w| w.intern }
 146     syms.unshift(nil)
 147     syms.push(nil)
 148
 149     syms.size.times { |i|
 150       [MAX_ORDER, syms.size-i].min.times { |ord|
 151         v = syms[i, ord+1]
 152         # puts "Learning #{v.inspect}"
 153         add(*v)
 154         # pp @mkv
 155       }
 156     }
 157   end
 158
 159   def learn(text, o={})
 160     opts = {:lowercase => true}.merge o
 161
 162     lc = opts[:lowercase]
 163
 164     simple_learn(text)
 165     if lc
 166       simple_learn(text.downcase)
 167     end
 168
 169     pp @mkv if defined? pp
 170   end
 171
 172 end
 173
 174 mkv = MarkovChainer.new
 175
 176 mkv.learn("This is a test, a nice little test indeed.")
 177