wookiee-jk commited on
Commit
74e8e71
โ€ข
1 Parent(s): 4977470
Files changed (1) hide show
  1. 12 +22 -0
12 ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import PreTrainedTokenizerFast, BartForConditionalGeneration
3
+
4
+ model_name = "ainize/kobart-news"
5
+ tokenizer = PreTrainedTokenizerFast.from_pretrained(model_name)
6
+ model = BartForConditionalGeneration.from_pretrained(model_name)
7
+
8
+ def summ(txt):
9
+ input_ids = tokenizer.encode(txt, return_tensors="pt")
10
+ summary_text_ids = model.generate(
11
+ input_ids = input_ids,
12
+ bos_token_id=model.config.bos_token_id, # BOS๋Š” Beginning of Sentence
13
+ eos_token_id=model.config.eos_token_id, # EOS๋Š” End of Sentence
14
+ length_penalty=2.0, # ์š”์•ฝ์„ ์–ผ๋งˆ๋‚˜ ์งง๊ฒŒ ํ• ์ง€
15
+ max_length=142,
16
+ min_length=56,
17
+ num_beams=4, # beam search
18
+ )
19
+ return tokenizer.decode(summary_text_ids[0], skip_special_tokens=True)
20
+
21
+ interface = gr.interface(summ, [gr.Textbox(label="original_text")], [gr.Textbox(label="summary")])
22
+ interface.launch()