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 MarkovChainer
  16   # Maximum depth
  17   MAX_ORDER = 5
  18
  19   # Word or nonword regexp:
  20   # can be used to scan a string splitting it into
  21   # words and nonwords.
  22   WNW = /\w+|\W/u
  23
  24   def initialize
  25     # mkv[i] holds the chains of order i
  26     @mkv = Array.new
  27
  28     MAX_ORDER.times { |i|
  29       @mkv[i] = {}
  30     }
  31
  32     # Each chain is in the form
  33     # [:array, :of, :symbols] => {
  34     #   :word => [chance before, chance after]
  35     #   :word => [chance before, chance after]
  36     # }
  37     # except for order 0, which is just a hash of
  38     # {:word => chance}
  39   end
  40
  41   def add_one(sym)
  42     s = sym.to_sym rescue nil
  43     if @mkv[0].has_key?(s)
  44       @mkv[0][s] += 1
  45     else
  46       @mkv[0][s] = 1
  47     end
  48   end
  49
  50   def add_before(array, prev)
  51     raise "Not enough words in new data" if array.empty?
  52     raise "Too many words in new data" if array.size > MAX_ORDER
  53     size = array.size
  54     if @mkv[size].has_key?(array)
  55       h = @mkv[size][array]
  56       if h.has_key?(prev)
  57        h[prev][0] += 1
  58       else
  59        h[prev] = [1,0]
  60       end
  61     else
  62       @mkv[size][array.dup] = { prev => [1, 0] }
  63     end
  64   end
  65
  66   def add_after(array, nxt)
  67     raise "Not enough words in new data" if array.empty?
  68     raise "Too many words in new data" if array.size > MAX_ORDER
  69     size = array.size
  70     if @mkv[size].has_key?(array)
  71       h = @mkv[size][array]
  72       if h.has_key?(nxt)
  73        h[nxt][1] += 1
  74       else
  75        h[nxt] = [0,1]
  76       end
  77     else
  78       @mkv[size][array.dup] = { nxt => [0, 1] }
  79     end
  80   end
  81
  82   def add_multi(array)
  83     raise "Too many words in new data" if array.size > MAX_ORDER + 1
  84     add_before(array.butfirst, array.first)
  85     add_after(array.butlast, array.last)
  86   end
  87
  88   def add(*data)
  89     if data.size == 1
  90       add_one(data.first)
  91     else
  92       add_multi(data)
  93     end
  94   end
  95
  96   def learn(text)
  97     syms = text.scan(WNW).map { |w| w.intern }
  98     syms.unshift(nil)
  99     syms.push(nil)
 100
 101     syms.size.times { |i|
 102       [MAX_ORDER, syms.size-i].min.times { |ord|
 103         v = syms[i, ord+1]
 104         # puts "Learning #{v.inspect}"
 105         add(*v)
 106         # pp @mkv
 107       }
 108     }
 109     pp @mkv if defined? pp
 110   end
 111
 112 end
 113
 114 mkv = MarkovChainer.new
 115
 116 mkv.learn("This is a test, this is a nice little test.")
 117