import streamlit as st
from Zmaker import Zmaker

if __name__ == "__main__":

    #ファインチューニング済みモデルの読み込み
    with st.spinner(text = "loading GPT-2..."):
        if not ("AI" in st.session_state.keys()):
            st.session_state["AI"] = Zmaker(
                ft_path = "./models/"
            )
    
    #設定用サイドバーの設定
    with st.sidebar:
        st.title("GPT-2のパラメータ")

        #max_lenの設定用スライダ
        sld_max_len = st.sidebar.slider(
            "length of the sentence", min_value = 0, max_value = 256, 
            value = (25, 75), step = 1, key = "length"
        )

        #temperatureの設定用スライダ
        sld_temp = st.sidebar.slider(
            "temperature", min_value = 0.1, max_value = 1.5, 
            value = 0.1, step = 0.1, key = "temp"
        )
        
        #top_kの設定用スライダ
        sld_top_k = st.sidebar.slider(
            "top_k", min_value = 0, max_value = 500, 
            value = 40, step = 1, key = "top_k"
        )
        
        #top_pの設定用スライダ
        sld_top_p = st.sidebar.slider(
            "top_p", min_value = 0.01, max_value = 1.0, 
            value = 0.95, step = 0.01, key = "top_p"
        )

        #repeat_ngram_sizeの設定用スライダ
        sld_top_p = st.sidebar.slider(
            "repeat_ngram_size ", min_value = 1, max_value = 10, 
            value = 1, step = 1, key = "repeat_ngram_size"
        )

    #メインフォームの設定
    with st.form(key = "Letter Form", clear_on_submit = False):
        st.title("おてがみ 入力欄")
        body = st.empty()
        if ("letter_body" in st.session_state.keys()): 
            ret = body.text_area(
                label = "おてがみを途中まで漢字+ひらがなで書くと続きをＡＩが生成します。"\
                        "生成できるのはおてがみ本文のみです。"\
                        "生成には数十秒～数分の時間がかかります。何卒ご了承ください。",
                value = st.session_state["letter_body"]
            )
        else:
            ret = body.text_area(
                label = "おてがみを途中まで漢字+ひらがなで書くと続きをＡＩが生成します。"\
                        "生成できるのはおてがみ本文のみです。"\
                        "生成には数十秒～数分の時間がかかります。何卒ご了承ください。",
                value = "ズッポシ村へようこそ！"
            )
        sub = st.form_submit_button("Generate")
    
    #注意事項
    with st.expander("zupposhi_makerについて/注意事項・免責事項"):
        st.markdown(
            """
            zupposhi_makerは稲葉百万鉄氏の[「どうぶつの森e+実況プレイ」](https://www.nicovideo.jp/mylist/45062007)の名物である、
            「イナババ怪文書」を用いて日本語GPT-2をファインチューニングして作成された生成AIです。

            本アプリの作成方法はエンジニア向け情報共有サービスzennにおいて[記事](https://zenn.dev/koujimachi2023/articles/b6bde79e73dd1d)としてまとめております。
            技術的な疑問点などありましたら、こちらの記事を参照ください。
            
            教師データの作成には以下のサイトを利用しました。
            - 稲葉百万鉄氏の[「どうぶつの森e+実況プレイ」](https://www.nicovideo.jp/mylist/45062007)において作成された文章
            - mintmama氏の[「ズッポシむら手紙集」](https://www.nicovideo.jp/series/85494)
            
            また、rinna社の配布する事前学習済みモデル(rinna/japanese-gpt2-medium)を
            ベースのモデルとして用いました。
            
            本アプリは稲葉百万鉄氏の動画作品に対するファン作品です。
            著作権法 第三十条の四に基づき今回の生成AIの学習は合法と解釈しておりますが、
            本アプリを利用する用途は個人利用の範囲に限定したいと考えております。
            そのため、本アプリおよびその生成データの**商用利用は一切禁止**とさせていただきます。
            
            また、以下の行為についても禁止とさせていただきます。
            - フェイクニュース作成への利用
            - 倫理上問題のある行為への利用(差別的な文章の生成・公開など)
            - わいせつ物の頒布
            - その他犯罪行為・迷惑行為への利用  
            
            本モデルの利用により利用者に何らかの不都合や損害が発生したとしても、
            モデルやデータセットの作者は何らの責任を負うものではありません。
            
            ただし、当方はWebアプリの公開など初心者です。
            調べられる範囲で対応しているつもりですが、
            著作権やライセンス、許可・禁止事項についての間違いがある可能性は大いにあります。
            つきましては、何かお気づきの点がありましたら遠慮なくご指摘いただけると幸いです。
            """
        )

       
    #submitボタンが押された
    if sub == True:
        #predictに必要な条件をGUIで設定した値に更新
        st.session_state["AI"].min_len = st.session_state["length"][0]
        st.session_state["AI"].max_len = st.session_state["length"][-1]
        st.session_state["AI"].top_k = st.session_state["top_k"]
        st.session_state["AI"].top_p = st.session_state["top_p"]
        st.session_state["AI"].temp = st.session_state["temp"]
        st.session_state["AI"].repeat_ngram_size = st.session_state["repeat_ngram_size"]

        #AIによる予測を実行
        with st.spinner(text = "generating..."):
            prompt = ret
            text  = str(st.session_state["AI"].GenLetter("<s>"+prompt)[0])
            text = text.replace('<s>', '')
            text = text.replace('</s>', '')
            st.session_state["letter_body"] = text
            st.experimental_rerun()