Update app.py

app.py

@@ -1,14 +1,19 @@
 import os
 import gradio as gr
 import torch
+import logging
 from langchain.chains import ConversationalRetrievalChain
 from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_community.document_loaders import PyPDFLoader
+from langchain.document_loaders import PyMuPDFLoader  # ✅ More stable PDF loader
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import Chroma
 from langchain_community.llms import HuggingFacePipeline
 from transformers import pipeline
 
+# Setup Logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 # Set Hugging Face Cache Directory
 os.environ["HF_HOME"] = "/tmp/huggingface_cache"
 
@@ -22,74 +27,53 @@ llm_pipeline = None
 embeddings = None
 persist_directory = "/tmp/chroma_db"  # Storage for vector DB
 
-
 def init_llm():
     """Initialize LLM and Embeddings"""
     global llm_pipeline, embeddings
-
+    
     hf_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
     if not hf_token:
         raise ValueError("HUGGINGFACEHUB_API_TOKEN is not set in environment variables.")
-
-    model_id = "tiiuae/falcon-rw-1b"
-
+    
+    model_id = "tiiuae/falcon-rw-1b"  # ✅ Can switch to "tiiuae/falcon-rw-1b" for lighter model
     hf_pipeline = pipeline("text-generation", model=model_id, device=DEVICE)
     llm_pipeline = HuggingFacePipeline(pipeline=hf_pipeline)
-
+    
     embeddings = HuggingFaceEmbeddings(
         model_name="sentence-transformers/all-MiniLM-L6-v2",
         model_kwargs={"device": DEVICE}
     )
-
-
-import time
+    logger.info("✅ LLM and Embeddings Initialized Successfully!")
 
 def process_document(file):
+    """Process uploaded PDF and create a retriever"""
     global conversation_retrieval_chain
-
+    
     if not llm_pipeline or not embeddings:
         init_llm()
-
-    start_time = time.time()
-    print(f"📂 Uploading PDF: {file.name}")
-
+    
     try:
-        # ✅ Ensure file is saved correctly
-        file_path = os.path.join("/tmp/uploads", file.name)
-        with open(file_path, "wb") as f:
-            f.write(file.read())
-        print(f"✅ PDF saved at {file_path} in {time.time() - start_time:.2f}s")
-
-        # ✅ Load PDF
-        start_time = time.time()
-        loader = PyPDFLoader(file_path)
-        documents = loader.load()
-        print(f"✅ PDF loaded in {time.time() - start_time:.2f}s")
+        file_path = file.name  # ✅ Ensures correct file path is passed
+        logger.info(f"📂 Processing PDF: {file_path}")
 
-        # ✅ Split text
-        start_time = time.time()
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=512, chunk_overlap=50)
+        loader = PyMuPDFLoader(file_path)  # ✅ Alternative loader for stability
+        documents = loader.load()
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1024, chunk_overlap=64)
         texts = text_splitter.split_documents(documents)
-        print(f"✅ Text split in {time.time() - start_time:.2f}s")
 
-        # ✅ Create ChromaDB
-        start_time = time.time()
-        db = Chroma.from_documents(texts, embedding=embeddings, persist_directory="/tmp/chroma_db")
-        print(f"✅ ChromaDB created in {time.time() - start_time:.2f}s")
+        # Load or create ChromaDB
+        db = Chroma.from_documents(texts, embedding=embeddings, persist_directory=persist_directory)
+        retriever = db.as_retriever(search_type="similarity", search_kwargs={'k': 6})
 
-        # ✅ Create retrieval chain
         conversation_retrieval_chain = ConversationalRetrievalChain.from_llm(
-            llm=llm_pipeline, retriever=db.as_retriever()
+            llm=llm_pipeline, retriever=retriever
         )
-        print("✅ Document processing complete!")
-
+        logger.info("✅ PDF Processed Successfully!")
         return "📄 PDF uploaded and processed successfully! You can now ask questions."
     
     except Exception as e:
-        print(f"❌ Error processing PDF: {str(e)}")
-        return f"Error: {str(e)}"
-
-
+        logger.error(f"❌ Error processing PDF: {str(e)}")
+        return f"❌ Error processing PDF: {str(e)}"
 
 def process_prompt(prompt, chat_history_display):
     """Generate a response using the retrieval chain"""
@@ -102,24 +86,21 @@ def process_prompt(prompt, chat_history_display):
     answer = output["answer"]
 
     chat_history.append((prompt, answer))
-    
     return chat_history
 
-
 # Define Gradio UI
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("<h1 style='text-align: center;'>Personal Data Assistant</h1>")
-
+    
     with gr.Row():
         dark_mode = gr.Checkbox(label="🌙 Toggle light/dark mode")
     
-    with gr.Column():  # ✅ Replace `gr.Box()` with `gr.Column()`
+    with gr.Column():
         gr.Markdown("Hello there! I'm your friendly data assistant, ready to answer any questions regarding your data. Could you please upload a PDF file for me to analyze?")
         file_input = gr.File(label="Upload File")
         upload_button = gr.Button("���� Upload File")
-
+    
     status_output = gr.Textbox(label="Status", interactive=False)
-
     chat_history_display = gr.Chatbot(label="Chat History")
     
     with gr.Row():
@@ -134,4 +115,4 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
 # Launch Gradio App
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch(server_name="0.0.0.0", server_port=7860)  # ✅ Works in Hugging Face Spaces