lib/epitome/corpus.rb in epitome-0.3.0 vs lib/epitome/corpus.rb in epitome-0.3.1

@@ -26,11 +26,11 @@
       # The number of documents in the corpus
       @n_docs = @original_corpus.keys.size
       
     end
 
-    def summary(summary_length, threshold=0.2)
+    def summary(summary_length, threshold=0.2, output=:string)
       s = @clean_corpus.values.flatten
       # n is the number of sentences in the total corpus
       n = @clean_corpus.values.flatten.size
 
       # Vector of Similarity Degree for each sentence in the corpus
@@ -56,10 +56,16 @@
       results = power_method similarity_matrix, n, 0.85
 
       # Ugly sleight of hand to return a text based on results
       # <Array>Results => <Hash>Results => <String>ResultsText
       h = Hash[@sentences.zip(results)]
-      return h.sort_by {|k, v| v}.reverse.first(summary_length).to_h.keys.join(" ")
+      if output == :array
+        return h.sort_by {|k, v| v}.reverse.first(summary_length).to_h.keys
+      elsif output == :string
+        return h.sort_by {|k, v| v}.reverse.first(summary_length).to_h.keys.join(" ")
+      else
+        return "No return format specified."
+      end
     end
 
     private
     def clean(sentence_array)
       # Clean the sentences a bit to avoid unnecessary operations