import gradio as gr

# make function using import pip to install torch
import pip
pip.main(['install', 'torch'])
pip.main(['install', 'transformers'])

import torch
import transformers


# saved_model
def load_model(model_path):
    saved_data = torch.load(
        model_path,
        map_location="cpu"
    )

    bart_best = saved_data["model"]
    train_config = saved_data["config"]
    tokenizer = transformers.PreTrainedTokenizerFast.from_pretrained('gogamza/kobart-base-v1')

    ## Load weights.
    model = transformers.BartForConditionalGeneration.from_pretrained('gogamza/kobart-base-v1')
    model.load_state_dict(bart_best)

    return model, tokenizer


# main
def inference(prompt):
    model_path = "./kobart-model-logical.pth"

    model, tokenizer = load_model(
        model_path=model_path
        )

    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 값
    examples=[
        "강원도 양양군에서 신청한 설악산국립공원 케이블카 사업이 7가지 부대조건을 지키는 조건으로 2015년 8월 28일 환경부 산하 국립공원위원회의 심의를 통과하였다. 제113차 국립공원위원회에서 내려진 이 결정으로 설악 오색-끝청 하단을 연결하는 3,492km의 케이블카과 건립될 예정이며, 설악산 대청봉과 1.4km 떨어진 지점까지 케이블카를 이용하여 접근할 수 있게 되었다. 사업비는 총 460억이며, 국비 50%, 군비 35%, 도비 15%로 결정되었다. 이 사업은 총 세차례에 걸친 반려와 부결 등을 통해 기획에 수정되어가다 결국 사업 타당성과 적정성을 향상시켰다는 평가를 받으며 확정되었다. 그러나 많은 시민단체와 환경단체는 이와 같은 결정에도 불구하고 국립공원의 환경파괴를 이유로 여전히 이를 반대하고 있으며, ‘산으로 간 4대강 사업’이라는 비난과 함께 국립공원이 유원지화 되는 비극을 막자고 주장한다. 편의성과 경제성의 논리와 함께 환경보존의 논리가 부딪히는 설악산국립공원의 케이블카는 건설되어야 하는가."
    ]
    ).launch() # launch(share=True)를 설정하면 외부에서 접속 가능한 링크가 생성됨

demo.launch()