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 $swig_runtime_data_type_pointer2 = nil
   6 require 'sqlite3'
   7
   8 class Array
   9   def sql_group
  10     "("+self.join(',')+")"
  11   end
  12
  13   def butlast
  14     first(self.size-1)
  15   end
  16
  17   def butfirst
  18     last(self.size-1)
  19   end
  20
  21   def pick_one
  22     self[rand(self.size)]
  23   end
  24
  25   def pick_some(n)
  26     return nil if self.empty?
  27     i = rand(self.size)
  28     if n < 0
  29       count = rand(self.size-i)
  30     else
  31       count = rand([n, self.size-i].min)
  32     end
  33     self[i, count]
  34   end
  35
  36   def pick_with_chance(chance)
  37     pick = rand(chance)
  38     self.each { |el, cch|
  39       ch = cch.to_i
  40       return el if pick < ch
  41       pick -= ch
  42     }
  43     nil
  44   end
  45 end
  46
  47 class String
  48   def quoted
  49     "'" + SQLite3::Database.quote(self)+"'"
  50   end
  51 end
  52
  53 class NilClass
  54   def quoted
  55     'NULL'
  56   end
  57 end
  58
  59 class MarkovChainer
  60   # Word or nonword regexp:
  61   # can be used to scan a string splitting it into
  62   # words and nonwords.
  63   WNW = /\w+|\W/u
  64
  65   attr_reader :max_order
  66
  67   def order(i=0)
  68     "order#{i}"
  69   end
  70
  71   def word(i)
  72     "word#{i}"
  73   end
  74
  75   def initialize(db, ord=5)
  76     @db = SQLite3::Database.new(db)
  77     @db.synchronous=0
  78     @max_order = ord
  79
  80     @db.execute("create table if not exists #{order(0).quoted} (#{word(0).quoted} text unique not null primary key, 'chance' integer not null default 0)")
  81     1.upto(@max_order+1) do |i|
  82       cols = (0..i).map { |j| word(j).quoted}
  83
  84       cmd = "create table if not exists "
  85       cmd << order(i).quoted + " ("
  86       cmd << cols.map { |c| c+' text'}.join(',')
  87       cmd << ", 'chance' integer not null default 0"
  88       cmd << ", unique#{cols.sql_group})"
  89       @db.execute(cmd)
  90     end
  91   end
  92
  93   def words
  94     @db.execute("select word0 from order0")
  95   end
  96
  97   def num_words
  98     @db.get_first_value("select count(*) from order0").to_i
  99   end
 100
 101   def where_selector(words, o={})
 102     offset = o[:offset].to_i
 103     ar = []
 104     words.length.times do |i|
 105       if words[i]
 106         ar << word(i+offset) + "=" + words[i].quoted
 107       else
 108         ar << word(i+offset) + " ISNULL"
 109       end
 110     end
 111     "where #{ar.join(' and ')}"
 112   end
 113
 114   def grouped_selector(words, o={})
 115     offset= o[:offset].to_i
 116     wds = []
 117     cols = []
 118     words.length.times do |i|
 119       cols << word(i)
 120       wds << words[i].quoted
 121     end
 122     if o.key?(:chance)
 123       cols << "chance"
 124       wds << o[:chance].to_i
 125     end
 126     return [cols.sql_group, wds.sql_group]
 127   end
 128
 129   def add_one(sym)
 130     # Don't add nil to order 0
 131     return unless sym
 132     @db.transaction do |db|
 133       if db.get_first_value("select chance from order0 where word0=?1", sym)
 134         db.execute("update order0 set chance=chance+1 where word0=?1", sym)
 135       else
 136         db.execute("insert into order0 (word0, chance) values (?1, 1)", sym)
 137       end
 138     end
 139     # puts @db.execute("select * from order0 where word0=?1", sym).inspect
 140   end
 141
 142   def add_multi(array)
 143     raise "Too many words in new data" if array.size > @max_order + 1
 144     table = order(array.length-1).quoted
 145     cols, wds = grouped_selector(array, :chance => 1)
 146     where = where_selector(array)
 147     @db.transaction do |db|
 148       if db.get_first_value("select chance from #{table} " + where)
 149         db.execute("update #{table} set chance=chance+1 " + where)
 150       else
 151         db.execute("insert into #{table} #{cols} values #{wds}")
 152       end
 153     end
 154     # puts @db.execute("select * from #{table} " + where).inspect
 155   end
 156
 157   def add(*data)
 158     # puts "adding #{data.inspect}"
 159     if data.size == 1
 160       add_one(data.first)
 161     else
 162       add_multi(data)
 163     end
 164   end
 165
 166   def simple_learn(text)
 167     return if text.empty?
 168     syms = text.scan(WNW)
 169     syms.unshift(nil)
 170     syms.push(nil)
 171
 172     syms.each_index do |i|
 173       max_len = [@max_order+1, syms.size - i].min
 174       1.upto max_len do |len|
 175         v = syms[i, len]
 176         # puts "Learning #{v.inspect}"
 177         add(*v)
 178         # pp @mkv
 179       end
 180     end
 181   end
 182
 183   def learn(text, o={})
 184     opts = {:lowercase => false}.merge o
 185
 186     lc = opts[:lowercase]
 187
 188     simple_learn(text)
 189     if lc
 190       simple_learn(text.downcase)
 191     end
 192   end
 193
 194   def raw_next(syms, o={})
 195     max_order = o.fetch(:max_order, @max_order)
 196     if max_order > syms.length
 197       max_order = syms.length
 198     end
 199     ar = syms.last(max_order)
 200     # puts "raw_next #{max_order} #{ar.inspect}"
 201
 202     table = order(max_order)
 203     sel = word(max_order)
 204     where = where_selector(ar)
 205
 206     choices = @db.execute("select #{sel},chance from #{table} #{where}")
 207     unless choices.empty?
 208       sum = @db.get_first_value("select sum(chance) from #{table} #{where}").to_i
 209       return choices.pick_with_chance(sum)
 210     else
 211       raw_next(ar.butfirst, o)
 212     end
 213   end
 214
 215   def next(text, o={})
 216     syms = text.scan(WNW)
 217     raw_next(syms, o)
 218   end
 219
 220   def raw_prev(syms, o={})
 221     max_order = o.fetch(:max_order, @max_order)
 222     if max_order > syms.length
 223       max_order = syms.length
 224     end
 225     ar = syms.first(max_order)
 226     # puts "raw_prev #{max_order} #{ar.inspect}"
 227
 228     table = order(max_order)
 229     sel = word(0)
 230     where = where_selector(ar,:offset => 1)
 231
 232     choices = @db.execute("select #{sel}, chance from #{table} #{where}")
 233     unless choices.empty?
 234       sum = @db.get_first_value("select sum(chance) from #{table} #{where}").to_i
 235       return choices.pick_with_chance(sum)
 236     else
 237       raw_prev(ar.butlast, o)
 238     end
 239   end
 240
 241   def prev(text, o={})
 242     syms = text.scan(WNW)
 243     raw_prev(syms, o)
 244   end
 245
 246   def complete_prev(text, o={})
 247     syms = text.scan(WNW)
 248     prev = raw_prev(syms, o)
 249     while prev do
 250       syms.unshift(prev)
 251       prev = raw_prev(syms, o)
 252     end
 253     return syms.to_s
 254   end
 255
 256   def complete_next(text, o={})
 257     syms = text.scan(WNW)
 258     nxt = raw_next(syms, o)
 259     while nxt do
 260       syms.push(nxt)
 261       nxt = raw_next(syms, o)
 262     end
 263     return syms.to_s
 264   end
 265
 266   def complete(text, o={})
 267     txt = text
 268     choices = @db.execute("select word0,chance from order0")
 269     return String.new if choices.empty?
 270     sum = @db.get_first_value("select sum(chance) from order0").to_i
 271     while txt.empty? do
 272       txt = choices.pick_with_chance(sum)
 273     end
 274     syms = [txt]
 275     prev = raw_prev(syms, o)
 276     nxt = raw_next(syms, o)
 277     while nxt or prev do
 278       # puts syms.inspect, nxt.inspect, prev.inspect
 279       # Keep adding only on the side where we
 280       # didn't come across a nil already
 281       if prev
 282         syms.unshift(prev)
 283         prev = raw_prev(syms, o)
 284       end
 285       if nxt
 286         syms.push(nxt)
 287         nxt = raw_next(syms, o)
 288       end
 289     end
 290     return syms.to_s
 291   end
 292
 293 end
 294