test

app.py

@@ -0,0 +1,58 @@
+"""
+baseline_interactive.py
+"""
+import gradio as gr
+import torch
+from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
+from transformers import pipeline
+
+model_name = "momo/rsp-sum"
+model = MBartForConditionalGeneration.from_pretrained(model_name)
+tokenizer = MBart50TokenizerFast.from_pretrained(model_name, src_lang="ko_KR", tgt_lang="ko_KR")
+
+# prefix = "translate English to German: "
+
+def summarization(model_name, text):
+    summarizer = pipeline("summarization", model=model, tokenizer=tokenizer)
+    summarizer("An apple a day, keeps the doctor away", min_length=50, max_length=150)
+
+    for result in summarizer(text):
+        print(result)
+    return result
+
+if __name__ == '__main__':
+
+    #Create a gradio app with a button that calls predict()
+    app = gr.Interface(
+        fn=summarization,
+        inputs='text',
+        outputs='text', 
+        title="News Summary Generator",
+        description="News Summary Generator"
+        )
+    app.launch()
+
+# with torch.no_grad():
+#     while True:
+#         t = input("\nDocument: ")
+#         tokens = tokenizer(
+#             t,
+#             return_tensors="pt",
+#             truncation=True,
+#             padding=True,
+#             max_length=600
+#         )
+
+#         input_ids = tokens.input_ids.cuda()
+#         attention_mask = tokens.attention_mask.cuda()
+
+#         sample_output = model.generate(
+#             input_ids, 
+#             max_length=150, 
+#             num_beams=5, 
+#             early_stopping=True,
+#             no_repeat_ngram_size=8,
+#     )
+#         # print("token:" + str(input_ids.detach().cpu()))
+#         # print("token:" + tokenizer.convert_ids_to_tokens(str(input_ids.detach().cpu())))
+#         print("Summary: " + tokenizer.decode(sample_output[0], skip_special_tokens=True))