# pip install streamlit
# pip install streamlit-chat

import streamlit as st
from streamlit_chat import message
import pandas as pd
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import json

@st.cache(allow_output_mutation=True)
def cached_model():
    model = SentenceTransformer('jhgan/ko-sroberta-multitask')
    return model

@st.cache(allow_output_mutation=True)
def get_dataset():
    df = pd.read_csv('wellness_dataset.csv')
    df['embedding'] = df.embedding.apply(json.loads)
    return df

model = cached_model()
df = get_dataset()

st.header('심리상담 챗봇')

if 'generated' not in st.session_state:
    st.session_state['generated'] = []
    
if 'past' not in st.session_state:
    st.session_state['past'] = []
    
with st.form('form', clear_on_submit=True):
    user_input = st.text_input('당신: ', '')
    submitted = st.form_submit_button('전송')
    
if submitted and user_input:
    embedding = model.encode(user_input)
    df['유사도'] = df.embedding.map(lambda x: cosine_similarity([embedding], [x]).squeeze())
    answer = df.loc[df.유사도.idxmax()]
    
    st.session_state.past.append(user_input)
    st.session_state.generated.append(answer.챗봇)
    
for i in range(len(st.session_state.past)-1, -1, -1):
    message(st.session_state.generated[i], key=str(i) + '_bot')
    if len(st.session_state.generated) > i:
        message(st.session_state.past[i], is_user=True, key=str(i) + '_user')
