import torch
import streamlit as st
from transformers import GPT2Tokenizer, GPT2LMHeadModel, PreTrainedTokenizerFast
import numpy as np


model = GPT2LMHeadModel.from_pretrained("jason9693/soongsil-univ-gpt-v1")
tokenizer = PreTrainedTokenizerFast.from_pretrained("jason9693/soongsil-univ-gpt-v1")

category_map = {
    "숭실대 에타": "<unused5>",
    "모두의 연애": "<unused3>",
    "대학생 잡담방": "<unused4>"
}

st.markdown("""# University Community KoGPT2 : 숭실대 에브리타임봇

[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1p6DIxsesi3eJNPwFwvMw0MeM5LkSGoPW?usp=sharing)	[![contributions welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg?style=flat)](https://github.com/jason9693/UCK-GPT2/issues)	![GitHub](https://img.shields.io/github/license/jason9693/UCK-GPT2)

## 대학 커뮤니티 게시글 생성기

SKT-AI에서 공개한 [KoGPT2](https://github.com/SKT-AI/KoGPT2) 모델을 파인튜닝하여 대학 커뮤니티 게시글을 생성하는 모델입니다. 총 에브리타임, 캠퍼스픽 데이터 22만개를 이용해서 학습을 진행했으며, 학습에는 대략 **3일**정도 소요되었습니다.

* [GPT 논문 리뷰 링크](https://www.notion.so/Improve-Language-Understanding-by-Generative-Pre-Training-GPT-afb4b5ef6e984961ac022b700c152b6b)

## 시연하기
""")


seed = st.text_input("Seed", "조만식 기념관")
category = st.selectbox("Category", list(category_map.keys()))
go = st.button("Generate")


st.markdown("## 생성 결과")
if go:
    input_context = category_map[category] + seed
    input_ids = tokenizer(input_context, return_tensors="pt")
    outputs = model.generate(
        input_ids=input_ids["input_ids"],
        max_length=250,
        num_return_sequences=1,
        no_repeat_ngram_size=3,
        repetition_penalty=2.0,
        do_sample=True,
        use_cache=True,
        eos_token_id=tokenizer.eos_token_id
    )
    st.write(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0].replace("<unused2>", "\n"))