add new

bart_demo_gradio.py

@@ -0,0 +1,49 @@
+import gradio as gr
+
+import torch
+import transformers
+
+# saved_model
+def load_model(model_path, config):
+    saved_data = torch.load(
+        model_path,
+        map_location="cpu" if config.gpu_id < 0 else "cuda:%d" % config.gpu_id
+    )
+
+    bart_best = saved_data["model"]
+    train_config = saved_data["config"]
+    tokenizer = transformers.PreTrainedTokenizerFast.from_pretrained(config.pretrained_model_name)
+
+    ## Load weights.
+    model = transformers.BartForConditionalGeneration.from_pretrained(config.pretrained_model_name)
+    model.load_state_dict(bart_best)
+
+    return model, tokenizer
+
+
+# main
+def inference(prompt):
+
+    config = define_argparser()
+    model_path = config.model_fpath
+
+    model, tokenizer = load_model(
+        model_path=model_path, 
+        config=config
+        )
+
+    input_ids = tokenizer.encode(prompt)
+    input_ids = torch.tensor(input_ids)
+    input_ids = input_ids.unsqueeze(0)
+    output = model.generate(input_ids)
+    output = tokenizer.decode(output[0], skip_special_tokens=True)    
+
+    return output
+
+demo = gr.Interface(
+    fn=inference, 
+    inputs="text", 
+    outputs="text" #return 값
+    ).launch(share=True) # launch(share=True)를 설정하면 외부에서 접속 가능한 링크가 생성됨
+
+demo.launch()