add luotuo summary

app.py

@@ -2,7 +2,7 @@ import numpy
 import torch
 import gradio as gr
 
-from transformers import PegasusForConditionalGeneration, Text2TextGenerationPipeline
+from transformers import PegasusForConditionalGeneration, Text2TextGenerationPipeline, AutoModel, AutoTokenizer
 from article_extractor.tokenizers_pegasus import PegasusTokenizer
 from embed import Embed
 
@@ -12,6 +12,9 @@ from harvesttext import HarvestText
 from sentence_transformers import SentenceTransformer, util
 from LexRank import degree_centrality_scores
 
+from luotuo_util import DeviceMap
+from peft import get_peft_model, LoraConfig, TaskType
+
 
 class SummaryExtractor(object):
     def __init__(self):
@@ -24,6 +27,39 @@ class SummaryExtractor(object):
         print(content)
         return str(self.text2text_genr(content, min_length=20, do_sample=False, num_return_sequences=3)[0]["generated_text"])
 
+class Tuoling_6B_extractor(object):
+    def __init__(self):
+        torch.set_default_tensor_type(torch.cuda.HalfTensor)
+        self.tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
+        self.model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True, device_map=DeviceMap("ChatGLM").get())
+
+        # load fine-tuned pretrained model.
+        peft_path = "./luotuoC.pt"
+        peft_config = LoraConfig(task_type=TaskType.CAUSAL_LM, inference_mode=True, r=8, lora_alpha=32, lora_dropout=0.1)
+        self.model = get_peft_model(self.model, peft_config)
+        self.model.load_state_dict(torch.load(peft_path), strict=False)
+        torch.set_default_tensor_type(torch.cuda.FloatTensor)
+
+    @staticmethod
+    def format_example(example: dict) -> dict:
+        context = f"Instruction: {example['instruction']}\n"
+        if example.get("input"):
+            context += f"Input: {example['input']}\n"
+        context += "Answer: "
+        target = example["output"]
+        return {"context": context, "target": target}
+
+    def extract(self, instruction: str, input=None) -> str:
+        with torch.no_grad():
+            feature = Tuoling_6B_extractor.format_example(
+                {"instruction": "请帮我总结以下内容", "output": "", "input": f"{instruction}"}
+            )
+            input_text = feature["context"]
+            input_ids = self.tokenizer.encode(input_text, return_tensors="pt")
+            out = self.model.generate(input_ids=input_ids, max_length=2048, temperature=0)
+            answer = self.tokenizer.decode(out[0])
+            return answer.split('Answer:')[1]
+
 class LexRank(object):
     def __init__(self):
         self.model = SentenceTransformer('paraphrase-multilingual-mpnet-base-v2')
@@ -56,35 +92,28 @@ class LexRank(object):
 
 # ---===--- worker instances ---===---
 t_randeng = SummaryExtractor()
+t_tuoling = Tuoling_6B_extractor()
+
 embedder = Embed()
 lex = LexRank()
 
 
 def randeng_extract(content):
     sentences = lex.find_central(content)
+    return str(list(t_randeng.extract(sentence) for sentence in sentences))
+
+def tuoling_extract(content):
+    sentences = lex.find_central(content)
+    return str(list(t_tuoling.extract(sentence) for sentence in sentences))
+
+def similarity_check(query, doc):
+    doc_list = doc.split("\n")
 
-    num = 500
-    ptr = 0
-    for index, sentence in enumerate(sentences):
-        num -= len(sentence)
-        if num < 0 and index > 0:
-            ptr = index - 1
-            break
-        if num < 0 and index == 0:
-            ptr = index
-            break
-    print(">>>")
-    for ele in sentences[:ptr]:
-        print(ele)
-    return t_randeng.extract("".join(sentences[:ptr]))
-
-
-def similarity_check(inputs: list):
-    doc_list = inputs[1].split("\n")
-    doc_list.append(inputs[0])
-    embedding_list = embedder.encode(doc_list)
-    scores = (embedding_list[-1] @ tf.transpose(embedding_list[:-1]))[0].numpy().tolist()
-    return numpy.array2string(scores, separator=',')
+    query_embedding = embedder.encode(query)
+    doc_embedding = embedder.encode(doc_list)
+    scores = (query_embedding @ tf.transpose(doc_embedding))[0].numpy().tolist()
+    # scores = list(util.cos_sim(embedding_list[-1], doc_embedding) for doc_embedding in embedding_list[:-1])
+    return str(scores)
 
 with gr.Blocks() as app:
     gr.Markdown("从下面的标签选择测试模块 [摘要生成,相似度检测]")
@@ -92,10 +121,14 @@ with gr.Blocks() as app:
     #     text_input = gr.Textbox()
     #     text_output = gr.Textbox()
     #     text_button = gr.Button("生成摘要")
-    with gr.Tab("Randeng-Pegasus-523M"):
+    with gr.Tab("LexRank->Randeng-Pegasus-523M"):
         text_input_1 = gr.Textbox(label="请输入长文本:", max_lines=1000)
         text_output_1 = gr.Textbox(label="摘要文本")
         text_button_1 = gr.Button("生成摘要")
+    with gr.Tab("LexRank->Tuoling-6B-chatGLM"):
+        text_input = gr.Textbox(label="请输入长文本:", max_lines=1000)
+        text_output = gr.Textbox(label="摘要文本")
+        text_button = gr.Button("生成摘要")
     with gr.Tab("相似度检���"):
         with gr.Row():
             text_input_query = gr.Textbox(label="查询文本")
@@ -103,7 +136,7 @@ with gr.Blocks() as app:
         text_button_similarity = gr.Button("对比相似度")
         text_output_similarity = gr.Textbox()
 
-    # text_button.click(tuoling_extract, inputs=text_input, outputs=text_output)
+    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)