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   # Word or nonword regexp:
  87   # can be used to scan a string splitting it into
  88   # words and nonwords.
  89   WNW = /\w+|\W/u
  90
  91   attr_reader :max_order
  92   def initialize(ord=5)
  93     # mkv[i] holds the chains of order i
  94     @max_order = ord+1
  95     @mkv = Array.new(@max_order) { |i|
  96       if i==0
  97         ChanceHash.new
  98       else
  99         @mkv[i] = Hash.new { |hash, key|
 100           hash[key] = {:prev => ChanceHash.new, :next => ChanceHash.new}
 101         }
 102       end
 103     }
 104   end
 105
 106   def add_one(sym)
 107     s = sym.to_sym rescue nil
 108     @mkv[0].increase(s)
 109   end
 110
 111   def add_before(array, prev)
 112     raise "Not enough words in new data" if array.empty?
 113     raise "Too many words in new data" if array.size > @max_order
 114     size = array.size
 115     h = @mkv[size][array.dup]
 116     h[:prev].increase(prev)
 117   end
 118
 119   def add_after(array, nxt)
 120     raise "Not enough words in new data" if array.empty?
 121     raise "Too many words in new data" if array.size > @max_order
 122     size = array.size
 123     h = @mkv[size][array.dup]
 124     h[:next].increase(nxt)
 125   end
 126
 127   def add_multi(array)
 128     raise "Too many words in new data" if array.size > @max_order + 1
 129     add_before(array.butfirst, array.first)
 130     add_after(array.butlast, array.last)
 131   end
 132
 133   def add(*data)
 134     if data.size == 1
 135       add_one(data.first)
 136     else
 137       add_multi(data)
 138     end
 139   end
 140
 141   def simple_learn(text)
 142     syms = text.scan(WNW).map { |w| w.intern }
 143     syms.unshift(nil)
 144     syms.push(nil)
 145
 146     syms.size.times { |i|
 147       ([@max_order, syms.size-i].min+1).times { |ord|
 148         v = syms[i, ord+1]
 149         # puts "Learning #{v.inspect}"
 150         add(*v)
 151         # pp @mkv
 152       }
 153     }
 154   end
 155
 156   def learn(text, o={})
 157     opts = {:lowercase => true}.merge o
 158
 159     lc = opts[:lowercase]
 160
 161     simple_learn(text)
 162     if lc
 163       simple_learn(text.downcase)
 164     end
 165
 166     pp @mkv if defined? pp
 167   end
 168
 169   def raw_next(syms)
 170     ar = syms.last([@max_order, syms.size].min)
 171     ord = ar.size
 172     if ord == 0
 173       @mkv[0].random
 174     else
 175       if @mkv[ord].key?(ar)
 176         @mkv[ord][ar][:next].random
 177       else
 178         raw_next(ar.last(ord-1))
 179       end
 180     end
 181   end
 182
 183   def next(text)
 184     syms = text.scan(WNW).map { |w| w.intern }
 185     raw_next(syms)
 186   end
 187
 188   def raw_prev(syms)
 189     ar = syms.first([@max_order, syms.size].min)
 190     ord = ar.size
 191     if ord == 0
 192       @mkv[0].random
 193     else
 194       if @mkv[ord].key?(ar)
 195         @mkv[ord][ar][:prev].random
 196       else
 197         raw_prev(ar.first(ord-1))
 198       end
 199     end
 200   end
 201
 202   def prev(text)
 203     syms = text.scan(WNW).map { |w| w.intern }
 204     raw_prev(syms)
 205   end
 206
 207   def complete_prev(text)
 208     syms = text.scan(WNW).map { |w| w.intern }
 209     prev = raw_prev(syms)
 210     while prev do
 211       syms.unshift(prev)
 212       prev = raw_prev(syms)
 213     end
 214     return syms.to_s
 215   end
 216
 217   def complete_next(text)
 218     syms = text.scan(WNW).map { |w| w.intern }
 219     nxt = raw_next(syms)
 220     while nxt do
 221       syms.push(nxt)
 222       nxt = raw_next(syms)
 223     end
 224     return syms.to_s
 225   end
 226
 227   def complete(text)
 228     txt = text
 229     while txt.empty? do
 230       txt = @mkv[0].random.to_s
 231     end
 232     syms = txt.scan(WNW).map { |w| w.intern }
 233     prev = raw_prev(syms)
 234     nxt = raw_next(syms)
 235     while nxt or prev do
 236       # Keep adding only on the side where we
 237       # didn't come across a nil already
 238       if prev
 239         syms.unshift(prev)
 240         prev = raw_prev(syms)
 241       end
 242       if nxt
 243         syms.push(nxt)
 244         nxt = raw_next(syms)
 245       end
 246     end
 247     return syms.to_s
 248   end
 249
 250 end
 251