import streamlit as st

@st.cache(allow_output_mutation=True)
def get_pipe():
    from transformers import AutoTokenizer, AutoModelForCausalLM
    model_name = "heegyu/ajoublue-gpt2-medium-summarization"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    tokenizer.truncation_side = "right"
    model = AutoModelForCausalLM.from_pretrained(model_name)
    return model, tokenizer

def get_response(tokenizer, model, context):
    context = f"내용: {context}\n요약: "
    inputs = tokenizer(
        context, 
        truncation=True,
        max_length=800,
        return_tensors="pt")
    
    generation_args = dict(
        max_new_tokens=128,
        min_length=inputs["input_ids"].shape[1] + 5,
        eos_token_id=2,
        do_sample=False,
        num_beams=4,
        early_stopping=True
    )

    outputs = model.generate(**inputs, **generation_args)
    response = tokenizer.decode(outputs[0])
    print(context)
    print(response)
    response = response[len(context):].replace("</s>", "")

    return response

st.title("한국어 요약 모델 demo")

with st.spinner("loading model..."):
    model, tokenizer = get_pipe()

input_ = st.text_area("요약할 문장을 입력하세요", value="미국과 중국의 갈등이 곳곳에서 불거지고 있다.\n남중국해에선 무력 충돌마저 벌어질 기세다.\n이런 갈등을 틈타 '넥스트 차이나'를 꿈꾸는 나라들의 움직임도 빨라지고 있다.\n가장 눈에 띄는 곳은 베트남이다.\n몇 년 전부터 '넥스트 차이나'로 부상해 온 베트남은 미중 무역전쟁의 최대 수혜국 중 하나로 꼽히는 곳이다.\n  애플이 베트남으로 생산공장을 옮기고 있고 LG전자, 퀄컴 등 다른 글로벌 기업들도 가세하고 있다.\n무엇보다 활기를 맞은 건 의류산업이다.\n사우스차이나모닝포스트(SCMP)는 최근 보도에서 \"몇 개월 전까지만 해도 미국에 의류 제품을 가장 많이 수출하는 나라는 중국이었지만, 베트남이 빠르게 치고 올라오고 있다\"고 설명했다.\n중국 제품이 여전히 제일 잘 팔리고는 있지만 \"살아남기 위한 몸부림\"으로 가격을 크게 낮춘 탓이란 분석이다.\n베트남의 경쟁력은 점점 더 커질 것으로 보인다.\n유럽연합(EU)과 자유무역협정(8월 발효)을 맺은 데다, 외국인 직접 투자도 급속히 늘고 있어서다.\n이 나라의 최대 강점은 젊은 인구가 많고 미국과 우호적인 관계를 맺고 있단 점이다.\n  정치적으로도 비교적 안정돼 있다.\n  미중 갈등이 본격화한 지난해부터 \"베트남이 황금기를 맞고 있다\"(워싱턴포스트)는 분석이 나온 이유다.\n또 다른 아시아 국가 인도도 이 기회를 놓치지 않으려 한다.\n국경 문제로 중국과 갈등을 빚고 있는 인도는 거대한 자국 시장을 무기로 중국 기업들을 압박하는 한편, 이참에 '넥스트 차이나' 자리도 꿰차겠단 야심을 내보이고 있다.\n최근 한국은행은 보고서를 내고 \"미국과의 공조가 강화할수록 인도의 글로벌 생산기지 역할은 더욱 커질 것\"이라고 분석했다.\n  중국보다 낮은 인건비 등을 내세워 '세계의 공장'을 노리고 있단 얘기다.\n가장 이목을 끄는 건 글로벌 IT 기업들과의 협력이다.\n최근 구글이 인도에 100억 달러(약 12조 원) 투자 계획을 밝혔고, 아마존과 퀄컴, 페이스북도 인도에 투자를 진행하고 있다.\n중국과 가까운 아시아 국가들만 이런 것이 아니다.\n미국과 이웃한 중미 국가 멕시코는 미중 무역전쟁이 시작될 때부터 '넥스트 차이나'를 노려온 국가다.\n최근 로이터통신은 \"아이폰 등을 위탁 생산하는 대만 폭스콘과 페가트론이 멕시코에 공장을 짓는 일을 검토하고 있다\"고 단독 보도했다.\n폭스콘은 중국 청두, 선전 등에 공장을 두고 130만 명을 고용하고 있는 업체다.")
ok = st.button("요약하기")
if input_ is not None and ok and len(input_) > 0:
    with st.spinner("요약을 생성중입니다..."):
        response = get_response(tokenizer, model, input_)
        st.text("요약 결과")
        st.success(response)