Spaces:
Build error
Build error
File size: 2,237 Bytes
1870c14 235b9c1 1870c14 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
import torch
import gradio as gr
from transformers import PegasusForConditionalGeneration, Text2TextGenerationPipeline
from article_extractor.tokenizers_pegasus import PegasusTokenizer
class SummaryExtractor(object):
def __init__(self):
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.model = PegasusForConditionalGeneration.from_pretrained('IDEA-CCNL/Randeng-Pegasus-523M-Summary-Chinese').to(self.device)
self.tokenizer = PegasusTokenizer.from_pretrained("IDEA-CCNL/Randeng-Pegasus-523M-Summary-Chinese")
self.text2text_genr = Text2TextGenerationPipeline(self.model, self.tokenizer, device=self.device)
def extract(self, content: str, min=20, max=30) -> str:
return str(self.text2text_genr(content, do_sample=False, min_length=min, max_length=max, num_return_sequences=3)[0]["generated_text"])
t_randeng = SummaryExtractor()
def randeng_extract(content):
return t_randeng.extract(content)
def similarity_check(query: str, doc: str):
doc_list = doc.split("\n")
return "similarity result"
with gr.Blocks() as app:
gr.Markdown("从下面的标签选择不同的摘要模型, 在左侧输入原文")
# with gr.Tab("CamelBell-Chinese-LoRA"):
# text_input = gr.Textbox()
# text_output = gr.Textbox()
# text_button = gr.Button("生成摘要")
with gr.Tab("Randeng-Pegasus-523M"):
text_input_1 = gr.Textbox()
text_output_1 = gr.Textbox()
text_button_1 = gr.Button("生成摘要")
# with gr.Tab("Flip Image"):
# with gr.Row():
# image_input = gr.Image()
# image_output = gr.Image()
# image_button = gr.Button("Flip")
with gr.Tab("相似度检测"):
with gr.Row():
text_input_query = gr.Textbox()
text_input_doc = gr.Textbox()
text_button_similarity = gr.Button("对比相似度")
text_output_similarity = gr.Textbox()
# 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, outputs=text_input_doc)
app.launch()
|