lib/twitter_ebooks/model.rb in twitter_ebooks-2.2.3 vs lib/twitter_ebooks/model.rb in twitter_ebooks-2.2.4

@@ -2,10 +2,11 @@
 # encoding: utf-8
 
 require 'json'
 require 'set'
 require 'digest/md5'
+require 'csv'
 
 module Ebooks
   class Model
     attr_accessor :hash, :sentences, :mentions, :keywords
 
@@ -23,9 +24,14 @@
 
       if path.split('.')[-1] == "json"
         log "Reading json corpus from #{path}"
         lines = JSON.parse(content, symbolize_names: true).map do |tweet|
           tweet[:text]
+        end
+      elsif path.split('.')[-1] == "csv"
+        log "Reading CSV corpus from #{path}"
+        lines = CSV.read(path).drop(1).map do |tweet|
+          tweet[5]
         end
       else
         log "Reading plaintext corpus from #{path}"
         lines = content.split("\n")
       end