Spaces:
Configuration error
Configuration error
| import streamlit as st | |
| import tempfile | |
| from langchain_huggingface import ChatHuggingFace, HuggingFaceEmbeddings, HuggingFaceEndpoint | |
| from langchain_core.prompts import PromptTemplate | |
| from langchain_text_splitters import RecursiveCharacterTextSplitter | |
| from langchain.document_loaders import PyPDFLoader | |
| from langchain.vectorstores import FAISS | |
| from langchain_core.output_parsers import StrOutputParser | |
| from langchain_core.messages import AIMessage, SystemMessage, HumanMessage | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| # ----------------------- | |
| # Initialize models | |
| # ----------------------- | |
| llm = HuggingFaceEndpoint( | |
| model='meta-llama/Llama-3.1-8B-Instruct', | |
| task='text-generation' | |
| ) | |
| chat_model = ChatHuggingFace(llm=llm) | |
| embed_model = HuggingFaceEmbeddings(model_name='sentence-transformers/all-mpnet-base-v2') | |
| # ----------------------- | |
| # Streamlit UI | |
| # ----------------------- | |
| st.set_page_config(page_title="Talk with your Documents", layout="wide") | |
| st.sidebar.title("Settings") | |
| chunk_size = st.sidebar.slider("Chunk Size", min_value=100, max_value=1000, value=300, step=50) | |
| chunk_overlap = st.sidebar.slider("Chunk Overlap", min_value=0, max_value=200, value=25, step=5) | |
| uploaded_files = st.sidebar.file_uploader( | |
| "Upload PDF(s)", | |
| type=["pdf"], | |
| accept_multiple_files=True | |
| ) | |
| # ----------------------- | |
| # Initialize FAISS store in session | |
| # ----------------------- | |
| if "vector_store" not in st.session_state: | |
| if uploaded_files: | |
| all_docs = [] | |
| for uploaded_file in uploaded_files: | |
| with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file: | |
| tmp_file.write(uploaded_file.read()) | |
| tmp_file_path = tmp_file.name | |
| loader = PyPDFLoader(tmp_file_path) | |
| docs = loader.load() | |
| all_docs.extend(docs) | |
| splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap) | |
| splitted_docs = splitter.split_documents(all_docs) | |
| st.session_state.vector_store = FAISS.from_documents(splitted_docs, embedding=embed_model) | |
| # ----------------------- | |
| # Chat logic | |
| # ----------------------- | |
| st.title("📄 Talk with your Documents - RAG Based") | |
| if "messages" not in st.session_state: | |
| st.session_state.messages = [SystemMessage(content="You are a helpful assistant")] | |
| query = st.text_input("Ask a question about your uploaded documents:") | |
| if st.button("Ask"): | |
| if not uploaded_files: | |
| st.error("❌ Please upload at least one PDF file before asking a question!") | |
| elif not query.strip(): | |
| st.warning("⚠️ Please enter a question to ask.") | |
| else: | |
| st.session_state.messages.append(HumanMessage(content=query)) | |
| with st.spinner("Finding the answer..."): | |
| # Use vector store from session | |
| retriever = st.session_state.vector_store.as_retriever(search_kwargs={'k': 3}) | |
| context_chunks = retriever.invoke(query) | |
| final_context = " ".join([chunk.page_content for chunk in context_chunks]) | |
| # Build chat history string | |
| history_str = "\n".join([f"{type(m).__name__}: {m.content}" for m in st.session_state.messages]) | |
| # Prepare prompt | |
| prompt = PromptTemplate( | |
| template=""" | |
| You are a helpful assistant. | |
| Answer the user's question ONLY based on the context provided. | |
| If the answer is not in the context, reply: 'This is not included in your File.' | |
| Context: {context} | |
| Question: {query} | |
| Previous chat history: | |
| {history} | |
| """ | |
| ) | |
| final_prompt = prompt.format(context=final_context, query=query, history=history_str) | |
| # Get answer | |
| parser = StrOutputParser() | |
| result = chat_model.invoke(final_prompt) | |
| answer = result.content | |
| st.session_state.messages.append(AIMessage(content=answer)) | |
| # ----------------------- | |
| # Display recent chat | |
| # ----------------------- | |
| st.subheader("💬 Recent Chat") | |
| for msg in st.session_state.messages[1:]: | |
| if isinstance(msg, HumanMessage): | |
| st.markdown(f"**You:** {msg.content}") | |
| elif isinstance(msg, AIMessage): | |
| st.markdown(f"**AI:** {msg.content}") | |