Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from langchain_openai import ChatOpenAI | |
| from langchain.memory import ConversationBufferMemory | |
| from langchain.chains import ConversationalRetrievalChain | |
| from langchain_community.document_loaders import PyPDFLoader | |
| from langchain_community.embeddings import HuggingFaceEmbeddings | |
| from langchain_community.vectorstores import FAISS | |
| from langchain.text_splitter import RecursiveCharacterTextSplitter | |
| import os | |
| import tempfile | |
| from dotenv import load_dotenv | |
| # Load environment variables | |
| load_dotenv() | |
| # Global variable to store the chat chain and vector store | |
| chat_chain = None | |
| vectordb = None | |
| # For debugging | |
| import sys | |
| print("Python version:", sys.version) | |
| def create_chat_chain(file_path): | |
| """ | |
| Create a conversational chain for the chat interface. | |
| Args: | |
| file_path: Path to the uploaded PDF file | |
| Returns: | |
| tuple: (status_message, chain, vector_store) | |
| """ | |
| global vectordb | |
| try: | |
| # Load the PDF document | |
| print(f"Loading document: {file_path}") | |
| loader = PyPDFLoader(file_path) | |
| docs = loader.load() | |
| if not docs: | |
| return "Error: No content could be extracted from the PDF.", None, None | |
| print(f"Successfully loaded {len(docs)} pages from the document") | |
| # Split text into chunks | |
| print("Splitting documents into chunks...") | |
| text_splitter = RecursiveCharacterTextSplitter( | |
| chunk_size=1500, # Increased chunk size for better context | |
| chunk_overlap=300, # Increased overlap for better temporal context | |
| length_function=len | |
| ) | |
| split_documents = text_splitter.split_documents(docs) | |
| print(f"Split into {len(split_documents)} chunks") | |
| # Create embeddings | |
| print("Creating embeddings...") | |
| embeddings = HuggingFaceEmbeddings( | |
| model_name="sentence-transformers/all-mpnet-base-v2" | |
| ) | |
| # Create vector store | |
| print("Creating vector store...") | |
| vectordb = FAISS.from_documents(split_documents, embeddings) | |
| # Set up retriever with more context | |
| retriever = vectordb.as_retriever( | |
| search_type="mmr", | |
| search_kwargs={ | |
| "k": 6, # Increased number of documents to retrieve | |
| "fetch_k": 20 # Increased fetch size for better context | |
| } | |
| ) | |
| # Set up memory with more context | |
| memory = ConversationBufferMemory( | |
| memory_key="chat_history", | |
| return_messages=True, | |
| output_key="answer", | |
| input_key="question" | |
| ) | |
| # Initialize the language model | |
| print("Initializing language model...") | |
| llm = ChatOpenAI( | |
| model_name="gpt-3.5-turbo", | |
| temperature=0.2, | |
| max_tokens=1500 | |
| ) | |
| # Create the conversational chain with better configuration | |
| print("Creating conversation chain...") | |
| from langchain.prompts import PromptTemplate | |
| from langchain.chains.conversational_retrieval.prompts import CONDENSE_QUESTION_PROMPT | |
| # Create a proper prompt template | |
| prompt_template = """Use the provided documents to answer the question. | |
| Be specific and accurate with temporal references. | |
| If you cannot find the information in the documents, say so clearly. | |
| {context} | |
| Question: {question} | |
| Answer:""" | |
| PROMPT = PromptTemplate( | |
| template=prompt_template, | |
| input_variables=["context", "question"] | |
| ) | |
| # Create the chain with proper configuration | |
| chain = ConversationalRetrievalChain.from_llm( | |
| llm=llm, | |
| retriever=retriever, | |
| memory=memory, | |
| return_source_documents=True, | |
| combine_docs_chain_kwargs={ | |
| "prompt": PROMPT | |
| } | |
| ) | |
| print("\nChain configuration:") | |
| print(f"LLM model: gpt-3.5-turbo") # Hardcoded since we know the model | |
| print(f"Retriever config: {retriever.search_kwargs}") | |
| print(f"Memory type: {type(memory).__name__}") | |
| print(f"Prompt template: {prompt_template}") | |
| print("Chat chain created successfully!") | |
| return "Document processed successfully! You can now ask questions about the PDF.", chain, vectordb | |
| except Exception as e: | |
| error_msg = f"Error processing document: {str(e)}" | |
| print(error_msg) | |
| import traceback | |
| traceback.print_exc() | |
| return error_msg, None, None | |
| def process_user_input(user_message, chat_history): | |
| """ | |
| Process user input and generate a response using the chat chain. | |
| Args: | |
| user_message (str): User's message | |
| chat_history: Current chat history | |
| Returns: | |
| str: Generated response | |
| """ | |
| global chat_chain | |
| if chat_chain is None: | |
| return "Please upload a PDF file first." | |
| try: | |
| print(f"Processing question: {user_message}") | |
| print(f"Current chat history: {len(chat_history)} messages") | |
| # Get response from the chain | |
| response = chat_chain({"question": user_message, "chat_history": chat_history}) | |
| return response["answer"] | |
| except Exception as e: | |
| import traceback | |
| print("\nError details:") | |
| print(f"Error type: {type(e).__name__}") | |
| print(f"Error message: {str(e)}") | |
| traceback.print_exc() | |
| if isinstance(e, ValueError): | |
| error_msg = f"Value error: {str(e)}" | |
| else: | |
| error_msg = "An unexpected error occurred. Please try again." | |
| return error_msg | |
| def handle_message(user_message, chat_history): | |
| """ | |
| Handle incoming messages and update chat history. | |
| Args: | |
| user_message (str): User's message | |
| history: Current chat history | |
| Returns: | |
| tuple: Empty message and updated history | |
| """ | |
| if not user_message.strip(): | |
| return "", chat_history | |
| # Get response from the chain | |
| response = process_user_input(user_message, chat_history) | |
| # Update chat history | |
| chat_history.append((user_message, response)) | |
| return "", chat_history | |
| def clear_chat(): | |
| """ | |
| Clear the conversation history. | |
| Returns: | |
| list: Empty list to clear the chat interface | |
| """ | |
| global chat_chain | |
| if chat_chain is not None and hasattr(chat_chain, 'memory') and chat_chain.memory is not None: | |
| chat_chain.memory.clear() | |
| return [] | |
| def process_uploaded_file(file): | |
| """ | |
| Process the uploaded PDF file and initialize the chat chain. | |
| Args: | |
| file: Uploaded file object from Gradio | |
| Returns: | |
| str: Status message | |
| """ | |
| global chat_chain, vectordb | |
| if file is None: | |
| return "No file uploaded. Please upload a PDF file." | |
| try: | |
| # Get the file path from the Gradio file object | |
| file_path = file.name | |
| # Validate file type | |
| if not file_path.lower().endswith('.pdf'): | |
| return "Error: Only PDF files are supported. Please upload a PDF file." | |
| # Create the chat chain with the uploaded file | |
| status_msg, new_chain, new_vectordb = create_chat_chain(file_path) | |
| if new_chain is not None and new_vectordb is not None: | |
| chat_chain = new_chain | |
| vectordb = new_vectordb | |
| return status_msg | |
| except Exception as e: | |
| import traceback | |
| traceback.print_exc() | |
| return f"Error processing file: {str(e)}" | |
| # Create the Gradio interface | |
| with gr.Blocks(theme=gr.themes.Soft()) as demo: | |
| gr.Markdown(""" | |
| # RAGBot with Memory | |
| Upload a PDF file and ask questions about its content. | |
| """) | |
| # File upload component | |
| file_output = gr.File(label="Upload a PDF file") | |
| upload_btn = gr.Button("Process PDF") | |
| status = gr.Markdown("Please upload a PDF file to begin.") | |
| # Chat interface | |
| chatbot = gr.Chatbot(height=400, label="Chat History") | |
| msg = gr.Textbox(placeholder="Type your message here...", label="Your Message") | |
| # Action buttons | |
| with gr.Row(): | |
| send_btn = gr.Button("Send") | |
| clear_btn = gr.Button("Clear Conversation") | |
| # Example questions | |
| examples = [ | |
| "What is this document about?", | |
| "Can you summarize the key points?", | |
| "What are the main topics covered?" | |
| ] | |
| # Add examples | |
| with gr.Row(): | |
| for example in examples: | |
| gr.Examples( | |
| examples=[[example]], | |
| inputs=msg, | |
| label=f"Try: {example[:20]}..." | |
| ) | |
| # Event handlers | |
| upload_btn.click( | |
| process_uploaded_file, | |
| inputs=file_output, | |
| outputs=status | |
| ) | |
| def handle_click(user_message, history): | |
| return handle_message(user_message, history) | |
| msg.submit( | |
| handle_click, | |
| inputs=[msg, chatbot], | |
| outputs=[msg, chatbot] | |
| ) | |
| send_btn.click( | |
| handle_click, | |
| inputs=[msg, chatbot], | |
| outputs=[msg, chatbot] | |
| ) | |
| clear_btn.click( | |
| clear_chat, | |
| outputs=chatbot | |
| ) | |
| # Launch the interface | |
| if __name__ == "__main__": | |
| print("Starting RAGBot application...") | |
| demo.launch(share=False) | |