import os
import streamlit as st
from transformers import (
    pipeline
)

# os.environ['HUGGING_FACE_HUB_TOKEN']=st.secrets["read_key"]
os.environ['HUGGING_FACE_HUB_TOKEN']="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"

# bert
pipe=pipeline(
    model="rjx/rjxai-ahs-en-v1",
)

if 'type' not in st.session_state:
    st.session_state.type=""

if 'article_key' not in st.session_state:
    st.session_state.article=""

if 'result' not in st.session_state:
    st.session_state.result=""

def form_article_click():
    if st.session_state.article_key == "":
        st.warning("Need to enter content")
    else:
        article=st.session_state.article_key
        # 截取大于510字符的内容
        if len(article)>=2048:
            article=article[0: 2048]
        result = pipe(article)
        print(result)
        if result:
            if result[0]["label"]==1:
                st.session_state.type="Human writing"
            elif result[0]["label"]==0:
                st.session_state.type="AI writing"
        st.session_state.result=result


st.set_page_config(layout="wide") 
st.title("AI write or Human write v2")

col1, col2 = st.columns(2)

with col1:
   st.header("input")
   with st.form("article_form"):
        # 文本输入
        st.text_area(
            'content',
            height=270,
            key='article_key'
        )
        # 提交按钮
        st.form_submit_button(label='submit', on_click=form_article_click)

with col2:
   st.header("output")
   st.text_input('classification', key="type", disabled=True)