fix

app.py

@@ -1,12 +1,14 @@
 import math, torch, gradio as gr
 
 from lex_rank import LexRank
+from lex_rank_new_model import LexRankNewModel
 from sentence_transformers import SentenceTransformer, util
 
 
 # ---===--- instances ---===---
 embedder = SentenceTransformer('paraphrase-multilingual-mpnet-base-v2')
 lex = LexRank()
+lex_new_model = LexRankNewModel()
 
 
 # 摘要方法
@@ -19,6 +21,16 @@ def extract_handler(content):
     return output
 
 
+# 摘要方法
+def extract_handler_new_model(content):
+    summary_length = math.ceil(len(content) / 10)
+    sentences = lex_new_model.find_central(content, num=summary_length)
+    output = ""
+    for index, sentence in enumerate(sentences):
+        output += f"{index}: {sentence}\n"
+    return output
+
+
 # 相似度检测方法
 def similarity_search(queries, doc):
     doc_list = doc.split('\n')
@@ -43,10 +55,14 @@ def similarity_search(queries, doc):
 #  web ui
 with gr.Blocks() as app:
     gr.Markdown("从下面的标签选择测试模块 [摘要生成,相似度检测]")
-    with gr.Tab("LexRank"):
+    with gr.Tab("LexRank-mpnet"):
         text_input_1 = gr.Textbox(label="请输入长文本:", lines=10, max_lines=1000)
         text_button_1 = gr.Button("生成摘要")
         text_output_1 = gr.Textbox(label="摘要文本(长度设置为原文长度的1/10)", lines=10)
+    with gr.Tab("LexRank-distiluse"):
+        text_input_2 = gr.Textbox(label="请输入长文本:", lines=10, max_lines=1000)
+        text_button_2 = gr.Button("生成摘要")
+        text_output_2 = gr.Textbox(label="摘要文本(长度设置为原文长度的1/10)", lines=10)
     with gr.Tab("相似度检测"):
         with gr.Row():
             text_input_query = gr.Textbox(lines=10, label="查询文本")
@@ -55,6 +71,7 @@ with gr.Blocks() as app:
         text_output_similarity = gr.Textbox()
 
     text_button_1.click(extract_handler, inputs=text_input_1, outputs=text_output_1)
+    text_button_2.click(extract_handler_new_model, inputs=text_input_2, outputs=text_output_2)
     text_button_similarity.click(similarity_search, inputs=[text_input_query, text_input_doc], outputs=text_output_similarity)
 
 app.launch(

lex_rank_new_model.py

@@ -0,0 +1,44 @@
+import numpy, nltk
+nltk.download('punkt')
+
+
+from harvesttext import HarvestText
+from lex_rank_util import degree_centrality_scores
+from sentence_transformers import SentenceTransformer, util
+
+
+class LexRankNewModel(object):
+    def __init__(self):
+        self.model = SentenceTransformer('distiluse-base-multilingual-cased-v1')
+        self.ht = HarvestText()
+
+    def find_central(self, content: str, num=100):
+        if self.contains_chinese(content):
+            sentences = self.ht.cut_sentences(content)
+        else:
+            sentences = nltk.sent_tokenize(content)
+        embeddings = self.model.encode(sentences, convert_to_tensor=True).cpu()
+
+        # Compute the pair-wise cosine similarities
+        cos_scores = util.cos_sim(embeddings, embeddings).numpy()
+
+        # Compute the centrality for each sentence
+        centrality_scores = degree_centrality_scores(cos_scores, threshold=None)
+
+        # We argsort so that the first element is the sentence with the highest score
+        most_central_sentence_indices = numpy.argsort(-centrality_scores)
+
+        # num = 100
+        res = []
+        for index in most_central_sentence_indices:
+            if num < 0:
+                break
+            res.append(sentences[index])
+            num -= len(sentences[index])
+        return res
+
+    def contains_chinese(self, content: str):
+        for _char in content:
+            if '\u4e00' <= _char <= '\u9fa5':
+                return True
+            return False