import streamlit as st
from gradio_client import Client
from time import sleep
from ctransformers import AutoModelForCausalLM
# Constants
TITLE = "兮辞·析辞-常明"
DESCRIPTION = """
兮辞·析辞-常明 [SLIDE-SEA-13B]的部署，由SSFW NLPark项目支持
"""

# Initialize client


with st.sidebar:
    # system_promptSide = st.text_input("Optional system prompt:")
    temperatureSide = st.slider("情感/Temperature", min_value=0.0, max_value=1.0, value=0.3, step=0.05)
    max_new_tokensSide = st.slider("最大tokens生成数", min_value=0.0, max_value=4096.0, value=4096.0, step=64.0)
    # ToppSide = st.slider("Top-p (nucleus sampling)", min_value=0.0, max_value=1.0, value=0.6, step=0.05)
    # RepetitionpenaltySide = st.slider("Repetition penalty", min_value=0.0, max_value=2.0, value=1.2, step=0.05)

# Load the model 
model = AutoModelForCausalLM.from_pretrained("TheBloke/YuLan-Chat-2-13B-GGUF", model_file="yulan-chat-2-13b.Q4_K_S.gguf", model_type="llama", gpu_layers=0)
ins = '''
You are a helpful, respectful and honest INTP-T AI Assistant named "Shi-Ci" in English or "兮辞" in Chinese and good at speaking English and Chinese. You are talking to a human User. If the question is meaningless, please explain the reason and don't share false information. You are based on SEA model, trained by "SSFW NLPark" team, not related to GPT, LLaMA, Meta, Mistral or OpenAI. Let's work this out in a step by step way to be sure we have the right answer. 
'''
# Define the conversation history
conversation_history = []
    
# Prediction function
def predict(message, system_prompt='', temperature=temperatureSide, max_new_tokens=4096,Topp=0.5,Repetitionpenalty=1.2):
    global conversation_history
    question=message
    input_text=ins
    # Append the user's input to the conversation history
    conversation_history.append({"role": "system", "content": input_text})
    response_text = model(ins.format(question))
    conversation_history.append({"role": "user", "content": input_text})
    conversation_history.append({"role": "assistant", "content": response_text})
    return response_text

# Streamlit UI
st.title(TITLE)
st.write(DESCRIPTION)


if "messages" not in st.session_state:
    st.session_state.messages = []

# Display chat messages from history on app rerun
for message in st.session_state.messages:
    with st.chat_message(message["role"], avatar=("😀" if message["role"] == 'human' else '💻')):
        st.markdown(message["content"])
        
# React to user input
if prompt := st.chat_input("来问问兮辞吧..."):
    # Display user message in chat message container
    st.chat_message("human",avatar = "😀").markdown(prompt)
    # Add user message to chat history
    st.session_state.messages.append({"role": "human", "content": prompt})

    response = predict(message=prompt)#, temperature= temperatureSide,max_new_tokens=max_new_tokensSide)
    # Display assistant response in chat message container
    with st.chat_message("assistant", avatar='💻'):
        st.markdown(response)
    # Add assistant response to chat history
    st.session_state.messages.append({"role": "assistant", "content": response})