Updated train.rb to use parts of speech