Spaces:
Sleeping
Sleeping
File size: 5,226 Bytes
1c5d91d 70acb39 1c5d91d a39157f 1c5d91d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 |
import streamlit as st
import os
from utils import load_background_image, apply_style, configure_page, breaks, file_uploader, initialise_session_state
from mylogging import configure_logging, toggle_logging, display_logs
from collections_setup import initialize_chromadb, initialize_collection, update_collection
from runpod_setup import get_relevant_text, generate_answer, get_contextual_prompt
if __name__ == "__main__":
configure_page()
apply_style()
load_background_image()
initialise_session_state()
breaks(2)
st.write(
"""
**Welcome!**
This Streamlit app lets you try out Retrieval-Augmented Generation (RAG) with the Mistral-7B model.
RAG helps the model find relevant information from documents before answering — so you get more accurate and helpful responses!
With this app, you can:
- Upload multiple PDF or text files to build a contextual knowledge base,
- Ask custom questions based on your uploaded documents, and
- Generate informed responses using a lightweight, hosted LLM.
"""
) # **Note:** All uploaded files and generated embeddings are stored **in memory only** and will be **lost when the app is closed or restarted**. No data is persisted between sessions.
# # Disable Chroma telemetry
os.environ["CHROMA_TELEMETRY_ENABLED"] = "False"
# Initialize logger
logger, log_stream = configure_logging()
st.markdown(
"""
<style>
/* This targets the selectbox container */
div[data-baseweb="select"] {
max-width: 150px;
}
</style>
""",
unsafe_allow_html=True,
)
st.divider()
# ---- Logging Setup ----
use_logging = False
if use_logging:
logging_level = st.selectbox("Select logging level", ['INFO', 'DEBUG', 'WARNING'], index=2)
toggle_logging(logging_level, logger)
# ---- Vector Store Setup ----
# Initialize ChromaDB and collection
EMBEDDING_MODEL = "all-MiniLM-L6-v2"
client, embedding_func = initialize_chromadb(EMBEDDING_MODEL)
collection_name = "my_collection"
collection = initialize_collection(client, embedding_func, collection_name)
# Upload files
st.markdown(
'<h3>Upload Files</h3>',
unsafe_allow_html=True)
st.html("""
Uploaded files are processed to build a contextual knowledge base for the RAG model.<br>
When you submit a prompt, the model retrieves relevant information from these documents to generate responses.
""")
col1_, _, col2_ = st.columns([.4, .1, .5])
with col1_:
file_uploader()
# Get the current uploaded filenames
logger.debug(f"\n\t-- Currently uploaded files: {st.session_state.get('uploaded_files_name', 'None')}")
# Update collection with uploaded files
files_to_add_to_collection= [file_name for file_name in st.session_state.get("uploaded_files_name", []) if file_name not in st.session_state.get("collections_files_name", [])]
logger.debug(f"\n\t-- Files not in collection: {files_to_add_to_collection}")
if files_to_add_to_collection:
collection = update_collection(collection, files_to_add_to_collection)
# Update the session state
logger.debug(f"Collection count: {collection.count()}")
logger.debug(f"\n\t-- Collection data currently uploaded:")
data_head = collection.get(limit=5)
for i, (metadata, document) in enumerate(zip(data_head["metadatas"], data_head["documents"]), start=1):
logger.debug(f"Item {i}:")
logger.debug(f"Metadata: {metadata}")
logger.debug(f"Document: {document}")
logger.debug("-" * 40)
# ---- Response Generation ----
# Streamlit UI
st.divider()
col1, _, col2 = st.columns([.6, .01, 1])
with col1:
st.subheader("Enter your prompt")
query = st.text_area("", height=100)
generate_clicked = st.button("Generate Response")
if generate_clicked:
if query.strip():
# Get the number of available documents in ChromaDB
available_docs = collection.count()
if available_docs > 0:
# Ensure n_results doesn't exceed available_docs
n_results = min(2, available_docs)
relevant_text = get_relevant_text(collection, query=query, nresults=n_results)
else:
relevant_text = "" # No documents available, so no additional context
st.warning("No knowledge base available. Generating response based only on the prompt.")
logger.debug("\n\t-- Relevant text retrieved:")
logger.debug(relevant_text)
with st.spinner("Generating response..."):
context_query = get_contextual_prompt(query, relevant_text)
response = generate_answer(context_query, max_tokens=200)
with col2:
st.subheader("Response:")
st.text_area("", value=response, height=200)
else:
logger.debug("No query provided; skipping relevant text retrieval.")
st.warning("Please enter a prompt.")
if use_logging:
display_logs(log_stream)
|