48 lines
1.6 KiB
Python
48 lines
1.6 KiB
Python
import streamlit as st
|
|
from rag import RAG
|
|
import re
|
|
import logging
|
|
|
|
@st.cache_resource
|
|
def init_rag():
|
|
llm_model_path = '/Users/peportier/llm/a/a/zephyr-7b-beta.Q5_K_M.gguf'
|
|
embed_model_name = 'intfloat/multilingual-e5-large'
|
|
collection_name = 'cera'
|
|
chromadb_path = './chromadb'
|
|
rag = RAG(llm_model_path, embed_model_name, collection_name, chromadb_path)
|
|
return rag
|
|
|
|
rag = init_rag()
|
|
|
|
st.title("CERA Chat")
|
|
|
|
if st.sidebar.button("Nouvelle conversation"):
|
|
st.session_state.messages = []
|
|
rag.reset_history()
|
|
|
|
if "messages" not in st.session_state:
|
|
st.session_state.messages = []
|
|
|
|
for message in st.session_state.messages:
|
|
with st.chat_message(message["role"]):
|
|
st.markdown(message["content"])
|
|
|
|
if prompt := st.chat_input("Comment puis-je vous aider ?"):
|
|
st.session_state.messages.append({"role": "user", "content": prompt})
|
|
with st.chat_message("user"):
|
|
st.markdown(prompt)
|
|
|
|
with st.chat_message("assistant"):
|
|
message_placeholder = st.empty()
|
|
full_response = ""
|
|
for response in rag.chat(prompt):
|
|
full_response += response
|
|
message_placeholder.markdown(full_response + "▌")
|
|
message_placeholder.markdown(full_response)
|
|
# url_pattern = r"URL :\n(https?://[^\s]+)"
|
|
# urls = re.findall(url_pattern, rag.chat_history[-1]['assistant'])
|
|
# markdown_urls = "\n".join([f"- {url}" for url in urls])
|
|
# logging.info(f"URLs: \n{markdown_urls}")
|
|
# message_placeholder.markdown(markdown_urls)
|
|
|
|
st.session_state.messages.append({"role": "assistant", "content": full_response}) |