Support page index for sentences

Author: dgrachev28

Sources/TextRank/Sentence.swift

@@ -14,10 +14,12 @@ public struct Sentence: Hashable {
         words.count
     }
 
+    public let pageIndex: Int
     public let originalTextIndex: Int
 
-    public init(text: String, originalTextIndex: Int, additionalStopwords: [String] = [String]()) {
+    public init(text: String, originalTextIndex: Int, pageIndex: Int = 0, additionalStopwords: [String] = [String]()) {
         self.text = text
+        self.pageIndex = pageIndex
         self.originalTextIndex = originalTextIndex
         words = Sentence.removeStopWords(from: Sentence.clean(self.text),
                                          additionalStopwords: additionalStopwords)

Sources/TextRank/TextRank.swift

@@ -8,7 +8,7 @@
 import Foundation
 
 public class TextRank {
-    public var text: String {
+    public var pages: [String] = [String]() {
         didSet {
             textToSentences()
         }
@@ -25,26 +25,35 @@ public class TextRank {
     }
 
     public init() {
-        text = ""
+        pages = [""]
         graph = TextGraph(damping: graphDamping)
     }
 
+    public init(pages: [String]) {
+        self.pages = pages
+        graph = TextGraph(damping: graphDamping)
+        textToSentences()
+    }
+
     public init(text: String) {
-        self.text = text
+        self.pages = [text]
         graph = TextGraph(damping: graphDamping)
         textToSentences()
     }
 
     public init(text: String, summarizationFraction: Float = 0.2, graphDamping: Float = 0.85) {
-        self.text = text
+        self.pages = [text]
         self.summarizationFraction = summarizationFraction
         self.graphDamping = graphDamping
         graph = TextGraph(damping: graphDamping)
         textToSentences()
     }
 
     func textToSentences() {
-        sentences = TextRank.splitIntoSentences(text, additionalStopwords: stopwords).filter { $0.length > 0 }
+        sentences = [];
+        for (pageIndex, page) in pages.enumerated() {
+            sentences.append(contentsOf: TextRank.splitIntoSentences(page, pageIndex: pageIndex, additionalStopwords: stopwords).filter { $0.length > 0 })
+        }
     }
 }
 
@@ -87,7 +96,7 @@ extension TextRank {
     /// Split text into sentences.
     /// - Parameter text: Original text.
     /// - Returns: An array of sentences.
-    static func splitIntoSentences(_ text: String, additionalStopwords stopwords: [String] = [String]()) -> [Sentence] {
+    static func splitIntoSentences(_ text: String, pageIndex: Int = 0, additionalStopwords stopwords: [String] = [String]()) -> [Sentence] {
         if text.isEmpty { return [] }
 
         var x = [Sentence]()
@@ -96,6 +105,7 @@ extension TextRank {
                 x.append(
                     Sentence(text: substring.trimmingCharacters(in: .whitespacesAndNewlines),
                              originalTextIndex: x.count,
+                             pageIndex: pageIndex,
                              additionalStopwords: stopwords)
                 )
             }

Tests/TextRankTests/TextRankTests.swift

@@ -49,7 +49,7 @@ class TextRankTests: XCTestCase {
         XCTAssertEqual(textRank.graph.edges.count, 2)
 
         text = "Dog cat bird. Sheep dog cat peacock. Horse cow fish dog chicken."
-        textRank.text = text
+        textRank.pages = [text]
         textRank.buildGraph()
         XCTAssertEqual(textRank.graph.nodes.count, 3)
         XCTAssertEqual(textRank.graph.edges.count, 3)