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
  95     @mkv = Array.new(@max_order+1) { |i|
  96       if i==0
  97         ChanceHash.new
  98       else
  99         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, o={})
 170     max_order = o.fetch(:max_order, @max_order)
 171     ar = syms.last([max_order, syms.size].min)
 172     ord = ar.size
 173     if ord == 0
 174       @mkv[0].random
 175     else
 176       if @mkv[ord].key?(ar)
 177         @mkv[ord][ar][:next].random
 178       else
 179         raw_next(ar.last(ord-1), o)
 180       end
 181     end
 182   end
 183
 184   def next(text, o={})
 185     syms = text.scan(WNW).map { |w| w.intern }
 186     raw_next(syms, o)
 187   end
 188
 189   def raw_prev(syms, o={})
 190     max_order = o.fetch(:max_order, @max_order)
 191     ar = syms.first([max_order, syms.size].min)
 192     ord = ar.size
 193     if ord == 0
 194       @mkv[0].random
 195     else
 196       if @mkv[ord].key?(ar)
 197         @mkv[ord][ar][:prev].random
 198       else
 199         raw_prev(ar.first(ord-1), o)
 200       end
 201     end
 202   end
 203
 204   def prev(text, o={})
 205     syms = text.scan(WNW).map { |w| w.intern }
 206     raw_prev(syms, o)
 207   end
 208
 209   def complete_prev(text, o={})
 210     syms = text.scan(WNW).map { |w| w.intern }
 211     prev = raw_prev(syms, o)
 212     while prev do
 213       syms.unshift(prev)
 214       prev = raw_prev(syms, o)
 215     end
 216     return syms.to_s
 217   end
 218
 219   def complete_next(text, o={})
 220     syms = text.scan(WNW).map { |w| w.intern }
 221     nxt = raw_next(syms, o)
 222     while nxt do
 223       syms.push(nxt)
 224       nxt = raw_next(syms, o)
 225     end
 226     return syms.to_s
 227   end
 228
 229   def complete(text, o={})
 230     txt = text
 231     while txt.empty? do
 232       txt = @mkv[0].random.to_s
 233     end
 234     syms = txt.scan(WNW).map { |w| w.intern }
 235     prev = raw_prev(syms, o)
 236     nxt = raw_next(syms, o)
 237     while nxt or prev do
 238       # Keep adding only on the side where we
 239       # didn't come across a nil already
 240       if prev
 241         syms.unshift(prev)
 242         prev = raw_prev(syms, o)
 243       end
 244       if nxt
 245         syms.push(nxt)
 246         nxt = raw_next(syms, o)
 247       end
 248     end
 249     return syms.to_s
 250   end
 251
 252 end
 253