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

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