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 size
  40     @hash.size
  41   end
  42
  43   def increase(el)
  44     if @hash.key?(el)
  45       @hash[el] += 1
  46     else
  47       @hash[el] = 1
  48     end
  49     @valid_pick = false
  50     return @hash[el]
  51   end
  52
  53   def decrease(el)
  54     if @hash.key?(el)
  55       @hash[el] -= 1
  56       @hash.delete(el) if @hash[el] == 0
  57     end
  58     @valid_pick = false
  59     return @hash[el]
  60   end
  61
  62   def make_picker
  63     @picker.clear
  64     total = 0
  65     @hash.each { |el, ch|
  66       total += ch
  67       @picker << [total, el]
  68     }
  69     @total = total
  70     @valid_pick = true
  71   end
  72
  73   def random
  74     case @hash.size
  75     when 0
  76       return nil
  77     when 1
  78       return @hash.keys.first
  79     else
  80       make_picker unless @valid_pick
  81       pick = rand(@total)
  82       @picker.each { |ch, el|
  83         return el if pick < ch
  84       }
  85     end
  86   end
  87 end
  88
  89 class MarkovChainer
  90   # Word or nonword regexp:
  91   # can be used to scan a string splitting it into
  92   # words and nonwords.
  93   WNW = /\w+|\W/u
  94
  95   attr_reader :max_order
  96   def initialize(ord=5)
  97     @max_order = ord
  98     @mkv = Hash.new { |hash, key|
  99       hash[key] = {:prev => ChanceHash.new, :next => ChanceHash.new}
 100     }
 101     @mkv[nil] = ChanceHash.new
 102   end
 103
 104   def add_one(sym)
 105     s = sym.to_sym rescue nil
 106     @mkv[nil].increase(s)
 107   end
 108
 109   def add_before(array, prev)
 110     raise "Not enough words in new data" if array.empty?
 111     raise "Too many words in new data" if array.size > @max_order
 112     h = @mkv[array.dup]
 113     h[:prev].increase(prev)
 114   end
 115
 116   def add_after(array, nxt)
 117     raise "Not enough words in new data" if array.empty?
 118     raise "Too many words in new data" if array.size > @max_order
 119     h = @mkv[array.dup]
 120     h[:next].increase(nxt)
 121   end
 122
 123   def add_multi(array)
 124     raise "Too many words in new data" if array.size > @max_order + 1
 125     add_before(array.butfirst, array.first)
 126     add_after(array.butlast, array.last)
 127   end
 128
 129   def add(*data)
 130     if data.size == 1
 131       add_one(data.first)
 132     else
 133       add_multi(data)
 134     end
 135   end
 136
 137   def simple_learn(text)
 138     syms = text.scan(WNW).map { |w| w.intern }
 139     syms.unshift(nil)
 140     syms.push(nil)
 141
 142     syms.size.times { |i|
 143       ([@max_order, syms.size-i].min+1).times { |ord|
 144         v = syms[i, ord+1]
 145         # puts "Learning #{v.inspect}"
 146         add(*v)
 147         # pp @mkv
 148       }
 149     }
 150   end
 151
 152   def learn(text, o={})
 153     opts = {:lowercase => false}.merge o
 154
 155     lc = opts[:lowercase]
 156
 157     simple_learn(text)
 158     if lc
 159       simple_learn(text.downcase)
 160     end
 161
 162     pp @mkv if defined? pp
 163   end
 164
 165   def raw_next(syms, o={})
 166     max_order = o.fetch(:max_order, @max_order)
 167     ar = syms.last([max_order, syms.size].min)
 168     if @mkv.key?(ar)
 169       @mkv[ar][:next].random
 170     else
 171       raw_next(ar.butfirst, o)
 172     end
 173   end
 174
 175   def next(text, o={})
 176     syms = text.scan(WNW).map { |w| w.intern }
 177     raw_next(syms, o)
 178   end
 179
 180   def raw_prev(syms, o={})
 181     max_order = o.fetch(:max_order, @max_order)
 182     ar = syms.first([max_order, syms.size].min)
 183     if @mkv.key?(ar)
 184       @mkv[ar][:prev].random
 185     else
 186       raw_prev(ar.butlast, o)
 187     end
 188   end
 189
 190   def prev(text, o={})
 191     syms = text.scan(WNW).map { |w| w.intern }
 192     raw_prev(syms, o)
 193   end
 194
 195   def complete_prev(text, o={})
 196     syms = text.scan(WNW).map { |w| w.intern }
 197     prev = raw_prev(syms, o)
 198     while prev do
 199       syms.unshift(prev)
 200       prev = raw_prev(syms, o)
 201     end
 202     return syms.to_s
 203   end
 204
 205   def complete_next(text, o={})
 206     syms = text.scan(WNW).map { |w| w.intern }
 207     nxt = raw_next(syms, o)
 208     while nxt do
 209       syms.push(nxt)
 210       nxt = raw_next(syms, o)
 211     end
 212     return syms.to_s
 213   end
 214
 215   def complete(text, o={})
 216     txt = text
 217     while txt.empty? do
 218       txt = @mkv[nil].random.to_s
 219     end
 220     syms = txt.scan(WNW).map { |w| w.intern }
 221     prev = raw_prev(syms, o)
 222     nxt = raw_next(syms, o)
 223     while nxt or prev do
 224       # Keep adding only on the side where we
 225       # didn't come across a nil already
 226       if prev
 227         syms.unshift(prev)
 228         prev = raw_prev(syms, o)
 229       end
 230       if nxt
 231         syms.push(nxt)
 232         nxt = raw_next(syms, o)
 233       end
 234     end
 235     return syms.to_s
 236   end
 237
 238 end
 239