Rubyで書くMovieLensデータセットの変換

MovieLensと呼ばれる映画の批評データを使って映画の推薦を行ってみよう，という試みです．
今回は元データをこれまでに作ったプログラムで使えるように変換するところまでやります．

MovieLensデータセットの変換

データの形式をこれまでのcritics(いわゆるHashのHash)と同じようにします．
やりかたは色々ありますが，ここではCSVライブラリを使いました．

require 'csv'

module My
  class MovieLens
    def initialize(item_file, critic_file)
      @item_file = item_file
      @critic_file = critic_file
    end

    #
    # critics形式に変換したHashを返す
    #
    def to_hash
      movies = extract_movies
      extract_critics(movies)
    end

    #
    # 映画IDと映画名の組み合わせを得る
    #
    def extract_movies
      movies = {}
      # カラムセパレータが'|'なので注意
      CSV.open(@item_file, 'r', '|') do |item|
        id    = item[0].to_i
        title = item[1].to_s
        movies[id] = title
      end
      movies
    end

    #
    # ユーザIDと映画の評価値の組み合わせを得る
    #
    def extract_critics(movies)
      critics = Hash.new({})
      # カラムセパレータがタブなので注意
      CSV.open(@critic_file, 'r', "\t") do |critic|
        user   = critic[0].to_i
        title  = movies[critic[1].to_i]
        rating = critic[2].to_f

        critics[user] = critics[user].merge({ title => rating })
      end
      critics
    end
  end
end

実行結果

このデータセットを後で使えるように，Marshalでファイルに保存します．

movie_lens = My::MovieLens.new('u.item', 'u.data')

File.open('movie_lens.dump', 'w') do |file|
  Marshal.dump(movie_lens.to_hash, file)
end

いざ，実行!

% ./movie_lens.rb
% ls -la movie_lens.dump 
-rw-r--r--  1 user  staff  719626  8 25 17:20 movie_lens.dump

MovieLensデータセットのでき上がり．