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

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