Create app.py

app.py

@@ -0,0 +1,229 @@
+import gradio as gr
+import spaces
+import os
+import logging
+from langchain.document_loaders import PyPDFLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import Chroma
+from huggingface_hub import InferenceClient, get_token
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Set HF_HOME for caching Hugging Face assets in persistent storage
+os.environ["HF_HOME"] = "/data/.huggingface"
+os.makedirs(os.environ["HF_HOME"], exist_ok=True)
+
+# Define persistent storage directories
+DATA_DIR = "/data"  # Root persistent storage directory
+DOCS_DIR = os.path.join(DATA_DIR, "documents")  # Subdirectory for uploaded PDFs
+CHROMA_DIR = os.path.join(DATA_DIR, "chroma_db")  # Subdirectory for Chroma vector store
+
+# Create directories if they don't exist
+os.makedirs(DOCS_DIR, exist_ok=True)
+os.makedirs(CHROMA_DIR, exist_ok=True)
+
+# Initialize Cerebras InferenceClient
+try:
+    token = get_token()
+    if not token:
+        logger.error("HF_TOKEN is not set in Space secrets")
+        client = None
+    else:
+        client = InferenceClient(
+            model="meta-llama/Llama-4-Scout-17B-16E-Instruct",
+            provider="cerebras",
+            token=token
+        )
+        logger.info("InferenceClient initialized successfully")
+except Exception as e:
+    logger.error(f"Failed to initialize InferenceClient: {str(e)}")
+    client = None
+
+# Global variables for vector store
+vectorstore = None
+retriever = None
+
+@spaces.GPU(duration=180)  # Use ZeroGPU (H200) for embedding generation, 180s timeout
+def initialize_rag(file):
+    global vectorstore, retriever
+    try:
+        # Debug file object properties
+        logger.info(f"File object: {type(file)}, Attributes: {dir(file)}")
+        logger.info(f"File name: {file.name}")
+
+        # Validate file
+        if not file or not file.name:
+            logger.error("No file provided or invalid file name")
+            return "Error: No file provided or invalid file name"
+        
+        # Verify temporary file exists and is accessible
+        if not os.path.exists(file.name):
+            logger.error(f"Temporary file {file.name} does not exist")
+            return f"Error: Temporary file {file.name} does not exist"
+        
+        # Check temporary file size
+        file_size = os.path.getsize(file.name)
+        logger.info(f"Temporary file size: {file_size} bytes")
+        if file_size == 0:
+            logger.error("Uploaded file is empty")
+            return "Error: Uploaded file is empty"
+
+        # Save uploaded file to persistent storage
+        file_name = os.path.basename(file.name)
+        file_path = os.path.join(DOCS_DIR, file_name)
+        
+        # Check if file exists and its size
+        should_save = True
+        if os.path.exists(file_path):
+            existing_size = os.path.getsize(file_path)
+            logger.info(f"Existing file {file_name} size: {existing_size} bytes")
+            if existing_size == 0:
+                logger.warning(f"Existing file {file_name} is empty, will overwrite")
+            else:
+                logger.info(f"File {file_name} already exists and is not empty, skipping save")
+                should_save = False
+
+        if should_save:
+            try:
+                with open(file.name, "rb") as src_file:
+                    file_content = src_file.read()
+                logger.info(f"Read {len(file_content)} bytes from temporary file")
+                if not file_content:
+                    logger.error("File content is empty after reading")
+                    return "Error: File content is empty after reading"
+                with open(file_path, "wb") as dst_file:
+                    dst_file.write(file_content)
+                    dst_file.flush()  # Ensure write completes
+                # Verify written file
+                written_size = os.path.getsize(file_path)
+                logger.info(f"Saved {file_name} to {file_path}, size: {written_size} bytes")
+                if written_size == 0:
+                    logger.error(f"Failed to write {file_name}, file is empty")
+                    return f"Error: Failed to write {file_name}, file is empty"
+            except PermissionError as e:
+                logger.error(f"Permission error writing to {file_path}: {str(e)}")
+                return f"Error: Permission denied writing to {file_path}"
+            except Exception as e:
+                logger.error(f"Error writing file to {file_path}: {str(e)}")
+                return f"Error writing file: {str(e)}"
+
+        # Load and split document
+        try:
+            loader = PyPDFLoader(file_path)
+            documents = loader.load()
+            if not documents:
+                logger.error("No content loaded from PDF")
+                return "Error: No content loaded from PDF"
+            text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+            texts = text_splitter.split_documents(documents)
+        except Exception as e:
+            logger.error(f"Error loading PDF: {str(e)}")
+            return f"Error loading PDF: {str(e)}"
+
+        # Create or update embeddings and vector store
+        try:
+            logger.info("Initializing HuggingFaceEmbeddings")
+            embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+            logger.info("Creating Chroma vector store")
+            vectorstore = Chroma.from_documents(
+                texts, embeddings, persist_directory=CHROMA_DIR
+            )
+            vectorstore.persist()  # Save to persistent storage
+            retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
+            logger.info(f"Vector store created and persisted to {CHROMA_DIR}")
+            return f"Document '{file_name}' processed and saved to {DOCS_DIR}!"
+        except Exception as e:
+            logger.error(f"Error in embeddings or Chroma: {str(e)}")
+            return f"Error processing embeddings: {str(e)}"
+    except Exception as e:
+        logger.error(f"Error processing document: {str(e)}")
+        return f"Error processing document: {str(e)}"
+
+def query_documents(query, history, system_prompt, max_tokens, temperature):
+    global retriever, client
+    try:
+        if client is None:
+            logger.error("InferenceClient not initialized")
+            return history, "Error: InferenceClient not initialized. Check HF_TOKEN."
+        if retriever is None:
+            logger.error("No documents loaded")
+            return history, "Error: No documents loaded. Please upload a document first."
+
+        # Ensure history is a list of [user, assistant] lists
+        logger.info(f"History before processing: {history}")
+        if not isinstance(history, list):
+            logger.warning("History is not a list, resetting")
+            history = []
+        history = [[str(item[0]), str(item[1])] for item in history if isinstance(item, (list, tuple)) and len(item) == 2]
+
+        # Retrieve relevant documents
+        docs = retriever.get_relevant_documents(query)
+        context = "\n".join([doc.page_content for doc in docs])
+
+        # Call Cerebras inference
+        logger.info("Calling Cerebras inference")
+        response = client.chat_completion(
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": f"Context: {context}\n\nQuery: {query}"}
+            ],
+            max_tokens=int(max_tokens),
+            temperature=float(temperature),
+            stream=False
+        )
+        answer = response.choices[0].message.content
+        logger.info("Inference successful")
+
+        # Update chat history with list format
+        history.append([query, answer])
+        logger.info(f"History after append: {history}")
+        return history, history
+    except Exception as e:
+        logger.error(f"Error querying documents: {str(e)}")
+        return history, f"Error querying documents: {str(e)}"
+
+# Load existing vector store on startup
+try:
+    if os.path.exists(CHROMA_DIR):
+        logger.info("Loading existing vector store")
+        embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+        vectorstore = Chroma(persist_directory=CHROMA_DIR, embedding_function=embeddings)
+        retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
+        logger.info(f"Loaded vector store from {CHROMA_DIR}")
+except Exception as e:
+    logger.error(f"Error loading vector store: {str(e)}")
+
+with gr.Blocks() as demo:
+    gr.Markdown("# RAG chatbot w/persistent storage (works best with CPU Upgrade)")
+
+    # File upload
+    file_input = gr.File(label="Upload Document (PDF)", file_types=[".pdf"])
+    file_output = gr.Textbox(label="Upload Status")
+    file_input.upload(initialize_rag, file_input, file_output)
+
+    # Chat interface
+    chatbot = gr.Chatbot(label="Conversation")
+
+    # Query and parameters
+    with gr.Row():
+        query_input = gr.Textbox(label="Query", placeholder="Ask about the document...")
+        system_prompt = gr.Textbox(
+            label="System Prompt",
+            value="You are a helpful assistant answering questions based on the provided document context."
+        )
+        max_tokens = gr.Slider(label="Max Tokens", minimum=50, maximum=2000, value=500, step=50)
+        temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=1.0, value=0.7, step=0.1)
+
+    # Submit button
+    submit_btn = gr.Button("Send")
+    submit_btn.click(
+        query_documents,
+        inputs=[query_input, chatbot, system_prompt, max_tokens, temperature],
+        outputs=[gr.Chatbot(), gr.Textbox()]
+    )
+
+if __name__ == "__main__":
+    demo.launch()