add custom siblings

lex_rank.py

@@ -1,9 +1,11 @@
+import math
+
 import numpy, nltk
 nltk.download('punkt')
 
 
 from harvesttext import HarvestText
-from lex_rank_util import degree_centrality_scores
+from lex_rank_util import degree_centrality_scores, find_siblings
 from sentence_transformers import SentenceTransformer, util
 
 
@@ -12,7 +14,7 @@ class LexRank(object):
         self.model = SentenceTransformer('paraphrase-multilingual-mpnet-base-v2')
         self.ht = HarvestText()
 
-    def find_central(self, content: str, num=10):
+    def find_central(self, content: str, num=10, siblings=0):
         if self.contains_chinese(content):
             sentences = self.ht.cut_sentences(content)
         else:
@@ -33,7 +35,7 @@ class LexRank(object):
         for index in most_central_sentence_indices:
             if num < 0:
                 break
-            res.append(sentences[index])
+            res.append(find_siblings(sentences, index, siblings)[1])
             num -= 1
         return res
 
@@ -42,3 +44,5 @@ class LexRank(object):
             if '\u4e00' <= _char <= '\u9fa5':
                 return True
         return False
+
+

lex_rank_L12.py

@@ -3,7 +3,7 @@ nltk.download('punkt')
 
 
 from harvesttext import HarvestText
-from lex_rank_util import degree_centrality_scores
+from lex_rank_util import degree_centrality_scores, find_siblings
 from sentence_transformers import SentenceTransformer, util
 
 
@@ -12,7 +12,7 @@ class LexRankL12(object):
         self.model = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2')
         self.ht = HarvestText()
 
-    def find_central(self, content: str, num=10):
+    def find_central(self, content: str, num=10, siblings=0):
         if self.contains_chinese(content):
             sentences = self.ht.cut_sentences(content)
         else:
@@ -33,7 +33,7 @@ class LexRankL12(object):
         for index in most_central_sentence_indices:
             if num < 0:
                 break
-            res.append(sentences[index])
+            res.append(find_siblings(sentences, index, siblings)[1])
             num -= 1
         return res
 

lex_rank_distiluse_v1.py → lex_rank_text2vec_v1.py

@@ -3,16 +3,16 @@ nltk.download('punkt')
 
 
 from harvesttext import HarvestText
-from lex_rank_util import degree_centrality_scores
+from lex_rank_util import degree_centrality_scores, find_siblings
 from sentence_transformers import SentenceTransformer, util
 
 
-class LexRankDistiluseV1(object):
+class LexRankText2VecV1(object):
     def __init__(self):
-        self.model = SentenceTransformer('distiluse-base-multilingual-cased-v1')
+        self.model = SentenceTransformer('shibing624/text2vec-base-chinese-paraphrase')
         self.ht = HarvestText()
 
-    def find_central(self, content: str, num=10):
+    def find_central(self, content: str, num=10, siblings=0):
         if self.contains_chinese(content):
             sentences = self.ht.cut_sentences(content)
         else:
@@ -33,7 +33,7 @@ class LexRankDistiluseV1(object):
         for index in most_central_sentence_indices:
             if num < 0:
                 break
-            res.append(sentences[index])
+            res.append(find_siblings(sentences, index, siblings)[1])
             num -= 1
         return res
 

lex_rank_util.py

@@ -6,7 +6,7 @@ Source: https://github.com/crabcamp/lexrank/tree/dev
 import numpy as np
 from scipy.sparse.csgraph import connected_components
 from scipy.special import softmax
-import logging
+import logging, math
 
 logger = logging.getLogger(__name__)
 
@@ -121,4 +121,12 @@ def stationary_distribution(
     if normalized:
         distribution /= n_1
 
-    return distribution
+    return distribution
+
+
+def find_siblings(sentences: [str], idx: int, siblings: int) -> (int, str):
+    if not siblings < math.ceil(len(sentences) / 2):
+        return -1, "siblings too large, try some value smaller."
+    head = max(idx - siblings, 0)
+    tail = min(idx + siblings + 1, len(sentences))
+    return 0, "".join(sentences[head:tail])