【ベイズ】Naive Bayes（単純ベイズ）による文書分類のサンプルプログラム【Ruby】

「[を] 【ベイズ】Naive Bayes（単純ベイズ）による文書分類のサンプルプログラム【Perl】」で紹介されている数式を理解するのはなかなか難しいので、Perlで書かれたプログラムを1行1行理解しながらRubyに書き換えてみました。

mkdat4nb.rb

#!/usr/bin/env ruby

$cnt_wc = Hash.new{|hash, value| hash[value] = Hash.new(0)}
$cnt_c  = Hash.new(0)
$cnt_w  = 0

while line = gets
  next unless line =~ /^(.+?)\t(.+?)$/
  cat, ws = $1, $2
  words = ws.split(',')
  words.each {|ele|
    next unless ele =~ /^(.+):([^:]+)$/
    w, tf = $1, $2.to_i
    $cnt_wc[w][cat] += tf
    $cnt_c[cat] += tf
    $cnt_w += tf
  }
end

$cnt_c.keys.sort.each {|cat|
  print " #{cat}\t#{$cnt_c[cat]}\t#{$cnt_w}\n"
}

$cnt_wc.keys.sort.each {|w|
  print "#{w}\t#{$cnt_wc[w].sort{|a, b| b[1] <=> a[1]}.map{|e| e.join(':')}.join(',')}\n"
}

実行結果

$ ./mkdat4nb.rb t.c2w > t.nb

nb.rb

#!/usr/bin/env ruby
# -*- coding: utf-8 -*-

$min_freq = 0.1

$cnt_wc = Hash.new{|hash, value| hash[value] = Hash.new }
$cnt_c  = Hash.new(0)
$cnt_w  = 0

File::open(ARGV[0]) {|fh|
  while line = fh.gets
    if line =~ /^\s(.+?)\t(\d+)\t(\d+)$/
      $cnt_c[$1] = $2.to_i
      $cnt_w = $3.to_i
    elsif line =~ /^(.+?)\t(.+?)$/
      w, cs = $1, $2
      cats = cs.split(",")
      cats.each {|e|
        next unless e =~ /^(.+):(\d+)$/
        c, f = $1, $2.to_i
        $cnt_wc[w][c] = f
      }
    end
  end
}

File::open(ARGV[1]) {|fh|
  while line = fh.gets
    line.chomp!
    ws = line.split(/ +/)
    print "> #{ws.join(' ')}\n"
    val = Hash.new(0)
    c2w = Hash.new{|hash, value| hash[value] = Hash.new }
    ws.each {|w|
      next unless $cnt_wc[w]
      $cnt_c.keys.each {|c|
        wc = $cnt_wc[w][c] || $min_freq
        val[c] += -1 * Math.log(wc / $cnt_c[c].to_f)
        c2w[c][w] = wc if wc >= 1
      }
    }
    val.keys.each {|c|
      val[c] += -1 * Math.log($cnt_c[c] / $cnt_w.to_f)
    }
    val.sort{|a, b| a[1] <=> b[1]}.each {|c|
      next if c2w[c[0]].empty? 
      v = (val[c[0]] * 1000).to_i / 1000.0
      print "#{c[0]}\t#{v}\t#{c2w[c[0]].sort{|a, b| b[1] <=> a[1]}.map{|e| e.join(':')}.join(',')}\n"
    }
  end
}

実行結果

$ ./nb.rb t.nb t.txt
> 選挙 移籍 離婚 破綻
エンタメ 11.843 離婚:7,破綻:2,移籍:1,選挙:1
政治 18.544 選挙:6
スポーツ 21.299 移籍:2
> 環境 医療 エネルギー 速報 遺伝子
科学 14.093 エネルギー:5,医療:4,遺伝子:3,環境:2
政治 20.284 エネルギー:2,医療:1,環境:1
スポーツ 25.488 エネルギー:1,速報:1
エンタメ 26.988 速報:1
> 防衛 オスプレイ 外交 破綻
政治 13.246 外交:3,オスプレイ:2,防衛:2
エンタメ 20.697 破綻:2
> 速報 サッカー 日本代表 移籍
スポーツ 10.114 サッカー:9,日本代表:8,移籍:2,速報:1
エンタメ 19.087 速報:1,移籍:1