felix-weiland's picture
Update app.py
03bd750
import streamlit as st
import os
from streamlit_chat import message
from llama_index.readers.schema.base import Document
from llama_index import LLMPredictor, GPTSimpleVectorIndex, PromptHelper, GPTTreeIndex
from langchain import OpenAI
import functions as f
import pandas as pd
#f.set_api_key("openai_key.txt") # Not needed as it's set in the huggingface environment.
llm_predictor = f.config_llm_predictor()
base_index = {}
for application in os.listdir("indices_vector"):
if application == ".DS_Store":
continue
else:
name = application.split(".")[0]
base_index[name] = f.load_index(f"indices_vector/{application}")
base_index = {key: value for key, value in sorted(base_index.items())}
df = f.get_data()
#Creating the chatbot interface
st.title("Chat with your reviews")
application = st.selectbox("Choose application", options=list(base_index.keys()))
index = base_index[application]
data = df[df["application"] == application.lower()]
# Storing the chat
if "generated" not in st.session_state:
st.session_state["generated"] = []
if "past" not in st.session_state:
st.session_state["past"] = []
if 'chat_sent' not in st.session_state:
st.session_state.chat_sent = ''
tab1, tab2 = st.tabs(["Chat", "Reviews"])
with tab1:
chat_input = f.get_chat_input()
if chat_input:
output = f.generate_response(chat_input, index, llm_predictor)
output = str(output).strip()
# store the output
st.session_state.past.append(chat_input)
st.session_state.generated.append(output)
# Empty state so that chat input is not accidentally resent
del chat_input
st.session_state.chat_sent = ''
# Push things
if st.session_state["generated"]:
for i in range(len(st.session_state["generated"]) - 1, -1, -1):
message(st.session_state["generated"][i], key=str(i))
message(st.session_state["past"][i], is_user=True, key=str(i) + "_user")
with tab2:
output = f.get_search(data)
if len(output) > 0:
for i, info in output[:100].iterrows():
st.write(info["review"], info["rating"], info["date"].split(" ")[0])
st.write("______")