diff --git a/chat_with_pdf/chat_pdf_llama3.2.py b/chat_with_pdf/chat_pdf_llama3.2.py new file mode 100644 index 0000000..911d950 --- /dev/null +++ b/chat_with_pdf/chat_pdf_llama3.2.py @@ -0,0 +1,69 @@ +# Import necessary libraries +import os +import tempfile +import streamlit as st +from embedchain import App +import base64 +from streamlit_chat import message + +# Define the embedchain_bot function +def embedchain_bot(db_path): + return App.from_config( + config={ + "llm": {"provider": "ollama", "config": {"model": "llama3.2:latest", "max_tokens": 250, "temperature": 0.5, "stream": True, "base_url": 'http://localhost:11434'}}, + "vectordb": {"provider": "chroma", "config": {"dir": db_path}}, + "embedder": {"provider": "ollama", "config": {"model": "llama3.2:latest", "base_url": 'http://localhost:11434'}}, + } + ) + +# Add a function to display PDF +def display_pdf(file): + base64_pdf = base64.b64encode(file.read()).decode('utf-8') + pdf_display = f'' + st.markdown(pdf_display, unsafe_allow_html=True) + +st.title("Chat with PDF using Llama 3.2") +st.caption("This app allows you to chat with a PDF using Llama 3.2 running locally with Ollama!") + +# Define the database path +db_path = tempfile.mkdtemp() + +# Create a session state to store the app instance and chat history +if 'app' not in st.session_state: + st.session_state.app = embedchain_bot(db_path) +if 'messages' not in st.session_state: + st.session_state.messages = [] + +# Sidebar for PDF upload and preview +with st.sidebar: + st.header("PDF Upload") + pdf_file = st.file_uploader("Upload a PDF file", type="pdf") + + if pdf_file: + st.subheader("PDF Preview") + display_pdf(pdf_file) + + if st.button("Add to Knowledge Base"): + with st.spinner("Adding PDF to knowledge base..."): + with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as f: + f.write(pdf_file.getvalue()) + st.session_state.app.add(f.name, data_type="pdf_file") + os.remove(f.name) + st.success(f"Added {pdf_file.name} to knowledge base!") + +# Chat interface +for i, msg in enumerate(st.session_state.messages): + message(msg["content"], is_user=msg["role"] == "user", key=str(i)) + +if prompt := st.chat_input("Ask a question about the PDF"): + st.session_state.messages.append({"role": "user", "content": prompt}) + message(prompt, is_user=True) + + with st.spinner("Thinking..."): + response = st.session_state.app.chat(prompt) + st.session_state.messages.append({"role": "assistant", "content": response}) + message(response) + +# Clear chat history button +if st.button("Clear Chat History"): + st.session_state.messages = [] \ No newline at end of file