Update app.py

app.py

@@ -1,6 +1,11 @@
 import streamlit as st
 import requests
 import logging
+from langchain.document_loaders import PDFPlumberLoader
+from langchain.text_splitters import RecursiveCharacterTextSplitter
+from langchain.prompts import ChatPromptTemplate
+from langchain.llms import HuggingFacePipeline
+from transformers import pipeline
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -61,6 +66,36 @@ def query(payload, api_url):
         logger.error(f"Failed to decode JSON response: {response.text}")
         return None
 
+# Function to load and process PDF
+def process_pdf(uploaded_file):
+    loader = PDFPlumberLoader(uploaded_file)
+    documents = loader.load()
+
+    # Split the documents into chunks
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=1000,
+        chunk_overlap=200,
+        add_start_index=True
+    )
+    return text_splitter.split_documents(documents)
+
+# Function to generate response using LangChain
+def generate_response_with_langchain(question, context):
+    prompt_template = """
+    You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Use three sentences maximum and keep the answer concise.
+    Question: {question} 
+    Context: {context} 
+    Answer:
+    """
+
+    prompt = ChatPromptTemplate.from_template(prompt_template)
+    model = HuggingFacePipeline(pipeline("text-generation", model=selected_model))
+
+    # Use LangChain to generate an answer
+    chain = prompt | model
+    response = chain.invoke({"question": question, "context": context})
+    return response
+
 # Chat interface
 st.title("🤖 DeepSeek Chatbot")
 st.caption("Powered by Hugging Face Inference API - Configure in sidebar")
@@ -70,59 +105,75 @@ for message in st.session_state.messages:
     with st.chat_message(message["role"]):
         st.markdown(message["content"])
 
-# Handle input
-if prompt := st.chat_input("Type your message..."):
-    st.session_state.messages.append({"role": "user", "content": prompt})
-
-    with st.chat_message("user"):
-        st.markdown(prompt)
-
-    try:
-        with st.spinner("Generating response..."):
-            # Prepare the payload for the API
-            # Combine system message and user input into a single prompt
-            full_prompt = f"{system_message}\n\nUser: {prompt}\nAssistant:"
-            payload = {
-                "inputs": full_prompt,
-                "parameters": {
-                    "max_new_tokens": max_tokens,
-                    "temperature": temperature,
-                    "top_p": top_p,
-                    "return_full_text": False
+# Handle input and PDF processing
+uploaded_file = st.file_uploader("Upload PDF", type="pdf", accept_multiple_files=False)
+if uploaded_file:
+    documents = process_pdf(uploaded_file)
+    context = "\n\n".join([doc.page_content for doc in documents])
+
+    # Ask the user a question
+    if prompt := st.chat_input("Type your message..."):
+        st.session_state.messages.append({"role": "user", "content": prompt})
+
+        with st.chat_message("user"):
+            st.markdown(prompt)
+
+        try:
+            with st.spinner("Generating response..."):
+                # Combine system message and user input into a single prompt
+                full_prompt = f"{system_message}\n\nUser: {prompt}\nAssistant:"
+                payload = {
+                    "inputs": full_prompt,
+                    "parameters": {
+                        "max_new_tokens": max_tokens,
+                        "temperature": temperature,
+                        "top_p": top_p,
+                        "return_full_text": False
+                    }
                 }
-            }
 
-            # Dynamically construct the API URL based on the selected model
-            api_url = f"https://api-inference.huggingface.co/models/{selected_model}"
-            logger.info(f"Selected model: {selected_model}, API URL: {api_url}")
+                # Dynamically construct the API URL based on the selected model
+                api_url = f"https://api-inference.huggingface.co/models/{selected_model}"
+                logger.info(f"Selected model: {selected_model}, API URL: {api_url}")
 
-            # Query the Hugging Face API using the selected model
-            output = query(payload, api_url)
+                # Query the Hugging Face API using the selected model
+                output = query(payload, api_url)
 
-            # Handle API response
-            if output is not None and isinstance(output, list) and len(output) > 0:
-                if 'generated_text' in output[0]:
-                    # Extract the assistant's response
-                    assistant_response = output[0]['generated_text'].strip()
+                # Handle API response
+                if output is not None and isinstance(output, list) and len(output) > 0:
+                    if 'generated_text' in output[0]:
+                        assistant_response = output[0]['generated_text'].strip()
 
-                    # Check for and remove duplicate responses
-                    responses = assistant_response.split("\n</think>\n")
-                    unique_response = responses[0].strip()
+                        # Check for and remove duplicate responses
+                        responses = assistant_response.split("\n</think>\n")
+                        unique_response = responses[0].strip()
 
-                    logger.info(f"Generated response: {unique_response}")
+                        logger.info(f"Generated response: {unique_response}")
 
-                    # Append response to chat only once
-                    with st.chat_message("assistant"):
-                        st.markdown(unique_response)
+                        # Append response to chat only once
+                        with st.chat_message("assistant"):
+                            st.markdown(unique_response)
 
-                    st.session_state.messages.append({"role": "assistant", "content": unique_response})
+                        st.session_state.messages.append({"role": "assistant", "content": unique_response})
+                    else:
+                        logger.error(f"Unexpected API response structure: {output}")
+                        st.error("Error: Unexpected response from the model. Please try again.")
                 else:
-                    logger.error(f"Unexpected API response structure: {output}")
-                    st.error("Error: Unexpected response from the model. Please try again.")
-            else:
-                logger.error(f"Empty or invalid API response: {output}")
-                st.error("Error: Unable to generate a response. Please check the model and try again.")
-
-    except Exception as e:
-        logger.error(f"Application Error: {str(e)}", exc_info=True)
-        st.error(f"Application Error: {str(e)}")
+                    logger.error(f"Empty or invalid API response: {output}")
+                    st.error("Error: Unable to generate a response. Please check the model and try again.")
+
+        except Exception as e:
+            logger.error(f"Application Error: {str(e)}", exc_info=True)
+            st.error(f"Application Error: {str(e)}")
+
+# Allow user to ask a question based on extracted PDF content
+if prompt := st.chat_input("Ask a question about the PDF content"):
+    if documents:
+        context = "\n\n".join([doc.page_content for doc in documents])  # Get context from documents
+        answer = generate_response_with_langchain(prompt, context)
+        
+        # Show the answer from LangChain model
+        with st.chat_message("assistant"):
+            st.markdown(answer)
+
+        st.session_state.messages.append({"role": "assistant", "content": answer})