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     self[0,self.size-1]
   9   end
  10   def butfirst
  11     self[1,self.size]
  12   end
  13 end
  14
  15 class ChanceHash
  16
  17   def initialize
  18     @hash = Hash.new(0)
  19     @picker = {}
  20     @total = 0
  21     @valid_pick = false
  22   end
  23
  24   def increase(el)
  25     if @hash.key?(el)
  26       @hash[el] += 1
  27     else
  28       @hash[el] = 1
  29     end
  30     @valid_pick = false
  31     return @hash[el]
  32   end
  33
  34   def decrease(el)
  35     if @hash.key?(el)
  36       @hash[el] -= 1
  37       @hash.delete(el) if @hash[el] == 0
  38     end
  39     @valid_pick = false
  40     return @hash[el]
  41   end
  42
  43   def make_picker
  44     @picker.clear
  45     total = 0
  46     @hash.each { |el, ch|
  47       total += ch
  48       @picker[total] = el
  49     }
  50     @total = total
  51     @valid_pick = true
  52   end
  53
  54   def random
  55     case @hash.size
  56     when 0
  57       return nil
  58     when 1
  59       return @hash.keys.first
  60     else
  61       make_picker unless @valid_pick
  62       pick = rand(@total)
  63       @picker.each { |ch, el|
  64         return el if pick < ch
  65       }
  66     end
  67   end
  68 end
  69
  70 class MarkovChainer
  71   # Maximum depth
  72   MAX_ORDER = 5
  73
  74   # Word or nonword regexp:
  75   # can be used to scan a string splitting it into
  76   # words and nonwords.
  77   WNW = /\w+|\W/u
  78
  79   def initialize
  80     # mkv[i] holds the chains of order i
  81     @mkv = Array.new
  82
  83     # Each chain is in the form
  84     # [:array, :of, :symbols] => {
  85     #   :prev => ChanceHash,
  86     #   :next => ChanceHash
  87     # }
  88     # except for order 0, which is a simple ChanceHash
  89     # itself
  90     MAX_ORDER.times { |i|
  91       if i == 0
  92         @mkv[0] = ChanceHash.new
  93       else
  94         @mkv[i] = Hash.new { |hash, key|
  95           hash[key] = {:prev => ChanceHash.new, :next => ChanceHash.new}
  96         }
  97       end
  98     }
  99
 100   end
 101
 102   def add_one(sym)
 103     s = sym.to_sym rescue nil
 104     @mkv[0].increase(s)
 105   end
 106
 107   def add_before(array, prev)
 108     raise "Not enough words in new data" if array.empty?
 109     raise "Too many words in new data" if array.size > MAX_ORDER
 110     size = array.size
 111     h = @mkv[size][array.dup]
 112     h[:prev].increase(prev)
 113   end
 114
 115   def add_after(array, nxt)
 116     raise "Not enough words in new data" if array.empty?
 117     raise "Too many words in new data" if array.size > MAX_ORDER
 118     size = array.size
 119     h = @mkv[size][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.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 => true}.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 end
 166
 167 mkv = MarkovChainer.new
 168
 169 mkv.learn("This is a test, a nice little test indeed.")
 170