Update app.py

app.py

@@ -1,10 +1,17 @@
-from transformers import MT5ForConditionalGeneration, AutoTokenizer, Text2TextGenerationPipeline
+from transformers import MT5ForConditionalGeneration, AutoTokenizer, Text2TextGenerationPipeline, AutoModelForSeq2SeqLM
 import gradio as gr
+import re
 
+# 翻译任务设置
 trans_mdl = MT5ForConditionalGeneration.from_pretrained("K024/mt5-zh-ja-en-trimmed")
 trans_tokenizer = AutoTokenizer.from_pretrained("K024/mt5-zh-ja-en-trimmed")
 trans_pipe = Text2TextGenerationPipeline(model=trans_mdl, tokenizer=trans_tokenizer)
 
+# 摘要任务设置
+sum_mdl = AutoModelForSeq2SeqLM.from_pretrained("csebuetnlp/mT5_multilingual_XLSum")
+sum_tokenizer = AutoTokenizer.from_pretrained("csebuetnlp/mT5_multilingual_XLSum")
+
+
 def translation_job(job, text):
     # 设置翻译任务和提示语的映射
     job_key = ["中译日", "中译英", "日译中", "英译中", "日译英", "英译日"]
@@ -15,7 +22,34 @@ def translation_job(job, text):
     print(input)
     response = trans_pipe(input, max_length=100, num_beams=4)
     return response[0]['generated_text']
-    
+
+
+def sum_job(text):
+    # 去除源文本中的空格
+    WHITESPACE_HANDLER = lambda k: re.sub('\s+', ' ', re.sub('\n+', ' ', k.strip()))
+
+    input_ids = sum_tokenizer(
+        [WHITESPACE_HANDLER(text)],
+        return_tensors="pt",
+        padding="max_length",
+        truncation=True,
+        max_length=512
+    )["input_ids"]
+
+    output_ids = sum_mdl.generate(
+        input_ids=input_ids,
+        max_length=84,
+        no_repeat_ngram_size=2,
+        num_beams=4
+    )[0]
+
+    response = sum_tokenizer.decode(
+        output_ids,
+        skip_special_tokens=True,
+        clean_up_tokenization_spaces=False
+    )
+
+    return response
 
 
 with gr.Blocks() as app:
@@ -29,8 +63,16 @@ with gr.Blocks() as app:
         source_text = gr.Textbox(lines=1, label="翻译文本", placeholder="请输入要翻译的文本")
         trans_result = gr.Textbox(lines=1, label="翻译结果")
         trans_btn = gr.Button("翻译")
-
+        
+    # 多语言自动摘要任务
+    with gr.Tab("多语言自动摘要"):
+        article_text = gr.Textbox(lines=8, label="待总结文本", placeholder="请输入要进行摘要的文本")
+        sum_result = gr.Textbox(lines=2, label="摘要结果")
+        sum_btn = gr.Button("摘要")
+        
     trans_btn.click(translation_job, inputs=[job_name, source_text], outputs=trans_result)
+    sum_btn.click(sum_job, inputs=article_text, outputs=sum_result)
+
 app.launch()