Skip to content

Commit 423e40c

Browse files
committed
test: removal of stopwords from TextRank sentences
1 parent 980bf3e commit 423e40c

File tree

2 files changed

+41
-1
lines changed

2 files changed

+41
-1
lines changed

Sources/TextRank/TextRank.swift

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,11 @@ public class TextRank {
1818
public var sentences = [Sentence]()
1919
public var summarizationFraction: Float = 0.2
2020
public var graphDamping: Float = 0.85
21-
public var stopwords = [String]()
21+
public var stopwords = [String]() {
22+
didSet {
23+
textToSentences()
24+
}
25+
}
2226

2327
public init() {
2428
text = ""

Tests/TextRankTests/TextRankTests.swift

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,4 +87,40 @@ class TextRankTests: XCTestCase {
8787
XCTAssertTrue(filteredResults.count < results.results.count)
8888
XCTAssertTrue(filteredResults.count == 2)
8989
}
90+
91+
func testStopwordsAreRemoved() {
92+
// Given
93+
let text = "Here are some sentences dog cat. With intentional stopwords gator. And some words that are not."
94+
95+
// When
96+
let textRank = TextRank(text: text)
97+
98+
// Then
99+
XCTAssertEqual(textRank.sentences.count, 2)
100+
XCTAssertEqual(textRank.sentences[0].length, 3)
101+
XCTAssertEqual(textRank.sentences.filter { $0.originalTextIndex == 0 }[0].words,
102+
Set(["sentences", "dog", "cat"]))
103+
XCTAssertEqual(textRank.sentences.filter { $0.originalTextIndex == 1 }[0].words,
104+
Set(["intentional", "stopwords", "gator"]))
105+
XCTAssertEqual(textRank.sentences[1].length, 3)
106+
}
107+
108+
func testAdditionalStopwords() {
109+
// Given
110+
let text = "Here are some sentences dog cat. With intentional stopwords gator. And some words that are not."
111+
let additionalStopwords = ["dog", "gator"]
112+
113+
// When
114+
let textRank = TextRank(text: text)
115+
textRank.stopwords = additionalStopwords
116+
117+
// Then
118+
XCTAssertEqual(textRank.sentences.count, 2)
119+
XCTAssertEqual(textRank.sentences[0].length, 2)
120+
XCTAssertEqual(textRank.sentences.filter { $0.originalTextIndex == 0 }[0].words,
121+
Set(["sentences", "cat"]))
122+
XCTAssertEqual(textRank.sentences.filter { $0.originalTextIndex == 1 }[0].words,
123+
Set(["intentional", "stopwords"]))
124+
XCTAssertEqual(textRank.sentences[1].length, 2)
125+
}
90126
}

0 commit comments

Comments
 (0)