fix

app.py

@@ -66,7 +66,7 @@ class LexRank(object):
         self.ht = HarvestText()
     def find_central(self, content: str):
         sentences = self.ht.cut_sentences(content)
-        embeddings = self.model.encode(sentences, convert_to_tensor=True)
+        embeddings = self.model.encode(sentences, convert_to_tensor=True).cpu()
 
         # Compute the pair-wise cosine similarities
         cos_scores = util.cos_sim(embeddings, embeddings).numpy()
@@ -78,19 +78,20 @@ class LexRank(object):
         most_central_sentence_indices = numpy.argsort(-centrality_scores)
 
         num = 100
-        ptr = 0
+        res = []
         for index in most_central_sentence_indices:
-            num -= len(sentences[index])
-            if num < 0 and index > 0:
-                ptr = index + 1
+            if num < 0:
                 break
-        return list(sentences[index] for index in most_central_sentence_indices[0: ptr])
+            res.append(sentences[index])
+            num -= len(sentences[index])
+        return res
 
 # ---===--- worker instances ---===---
-t_randeng = SummaryExtractor()
+# t_randeng = SummaryExtractor()
 # t_tuoling = Tuoling_6B_extractor()
 
-embedder = Embed()
+# embedder = Embed()
+embedder = SentenceTransformer('paraphrase-multilingual-mpnet-base-v2')
 lex = LexRank()
 
 
@@ -99,9 +100,9 @@ def randeng_extract(content):
     output = "原文: \n"
     for index, sentence in enumerate(sentences):
         output += f"{index}: {sentence}\n"
-    output += "摘要:\n"
-    for index, sentence in enumerate(sentences):
-        output += f"{index}: {t_randeng.extract(sentence)}\n"
+    # output += "摘要:\n"
+    # for index, sentence in enumerate(sentences):
+    #     output += f"{index}: {t_randeng.extract(sentence)}\n"
     return output
 
 # def tuoling_extract(content):
@@ -117,10 +118,30 @@ def similarity_check(query, doc):
     # scores = list(util.cos_sim(embedding_list[-1], doc_embedding) for doc_embedding in embedding_list[:-1])
     return str(scores)
 
+def similarity_search(queries, doc):
+    doc_list = doc.split('\n')
+    query_list = queries.split('\n')
+
+    corpus_embeddings = embedder.encode(doc_list, convert_to_tensor=True)
+    top_k = min(5, len(doc_list))
+    output = ""
+    for query in query_list:
+        query_embedding = embedder.encode(query, convert_to_tensor=True)
+        # We use cosine-similarity and torch.topk to find the highest 5 scores
+        cos_scores = util.cos_sim(query_embedding, corpus_embeddings)[0]
+        top_results = torch.topk(cos_scores, k=top_k)
+        output += "\n\n======================\n\n"
+        output += f"Query: {query}"
+        output += "\nTop 5 most similar sentences in corpus:"
+        for score, idx in zip(top_results[0], top_results[1]):
+            output += f"{doc_list[idx]}(Score: {score})"
+    return output
+
+
 with gr.Blocks() as app:
     gr.Markdown("从下面的标签选择测试模块 [摘要生成,相似度检测]")
     with gr.Tab("LexRank->Randeng-Pegasus-523M"):
-        text_input_1 = gr.Textbox(label="请输入长文本:", max_lines=1000)
+        text_input_1 = gr.Textbox(label="请输入长文本:", lines=10, max_lines=1000)
         text_output_1 = gr.Textbox(label="摘要文本", lines=10)
         text_button_1 = gr.Button("生成摘要")
     # with gr.Tab("LexRank->Tuoling-6B-chatGLM"):
@@ -136,7 +157,7 @@ with gr.Blocks() as app:
 
     # text_button.click(tuoling_extract, inputs=text_input, outputs=text_output)
     text_button_1.click(randeng_extract, inputs=text_input_1, outputs=text_output_1)
-    text_button_similarity.click(similarity_check, inputs=[text_input_query, text_input_doc], outputs=text_output_similarity)
+    text_button_similarity.click(similarity_search, inputs=[text_input_query, text_input_doc], outputs=text_output_similarity)
 
 app.launch(
     share=True,