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