Update app.py

app.py

@@ -1,217 +1,133 @@
 import streamlit as st
-from langchain.chains import RetrievalQA
-from langchain.vectorstores import Milvus
-from langchain.embeddings import HuggingFaceEmbeddings
-from transformers import AutoTokenizer
-from langchain_groq import ChatGroq
 import os
-from docling.document_converter import DocumentConverter, PdfFormatOption
-from docling.datamodel.base_models import InputFormat
-from docling.datamodel.pipeline_options import PdfPipelineOptions
-from docling_core.transforms.chunker.hybrid_chunker import HybridChunker
-from docling_core.types.doc.document import TableItem
-from langchain_core.documents import Document
-import itertools
-from docling_core.types.doc.labels import DocItemLabel
-import google.generativeai as genai
 from PIL import Image
-import base64
-import io
+import google.generativeai as genai
+from utils.document_processing import process_pdf
+from utils.models import load_models
+from utils.rag import query_pipeline
 
-# Initialize components (similar to your notebook)
-@st.cache_resource
-def initialize_components():
-    # Initialize embeddings
-    embeddings_model_path = "ibm-granite/granite-embedding-30m-english"
-    embeddings_model = HuggingFaceEmbeddings(model_name=embeddings_model_path)
-    embeddings_tokenizer = AutoTokenizer.from_pretrained(embeddings_model_path)
-    
-    # Initialize language model
-    GROQ_API_KEY = "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"
-    model = ChatGroq(model_name="llama3-70b-8192", api_key=GROQ_API_KEY)
+# Configure the app
+st.set_page_config(
+    page_title="PDF RAG Pipeline",
+    page_icon="📄",
+    layout="wide"
+)
+
+# Initialize session state
+if 'models_loaded' not in st.session_state:
+    st.session_state.models_loaded = False
+if 'processed_docs' not in st.session_state:
+    st.session_state.processed_docs = None
+
+# Sidebar for configuration
+with st.sidebar:
+    st.title("Configuration")
     
-    # Initialize vision model
-    GOOGLE_API_KEY = "AIzaSyBTt66oOvxpLeYn41sR-KkjSYPK2vOAqkU"
-    genai.configure(api_key=GOOGLE_API_KEY)
-    vision_model = genai.GenerativeModel(model_name="gemini-1.5-flash")
+    # API keys
+    groq_api_key = st.text_input("Groq API Key", type="password")
+    google_api_key = st.text_input("Google API Key", type="password")
     
-    return embeddings_model, embeddings_tokenizer, model, vision_model
-
-def process_pdf(file_path, embeddings_tokenizer, vision_model):
-    # PDF processing (similar to your notebook)
-    pdf_pipeline_options = PdfPipelineOptions(
-        do_ocr=True,
-        generate_picture_images=True
+    # Model selection
+    embedding_model = st.selectbox(
+        "Embedding Model",
+        ["ibm-granite/granite-embedding-30m-english"],
+        index=0
     )
     
-    format_options = {
-        InputFormat.PDF: PdfFormatOption(pipeline_options=pdf_pipeline_options),
-    }
-    
-    converter = DocumentConverter(format_options=format_options)
-    sources = [file_path]
-    conversions = {
-        source: converter.convert(source=source).document for source in sources
-    }
+    llm_model = st.selectbox(
+        "LLM Model",
+        ["llama3-70b-8192"],
+        index=0
+    )
     
-    # Process text chunks
-    doc_id = 0
-    texts = []
+    # File upload
+    uploaded_file = st.file_uploader(
+        "Upload a PDF file", 
+        type=["pdf"],
+        accept_multiple_files=False
+    )
     
-    for source, docling_document in conversions.items():
-        chunker = HybridChunker(tokenizer=embeddings_tokenizer)
-        
-        for chunk in chunker.chunk(docling_document):
-            items = chunk.meta.doc_items
-            
-            if len(items) == 1 and isinstance(items[0], TableItem):
-                continue
+    if st.button("Initialize Models"):
+        with st.spinner("Loading models..."):
+            try:
+                # Load models
+                embeddings_model, embeddings_tokenizer, vision_model, llm_model = load_models(
+                    embedding_model=embedding_model,
+                    llm_model=llm_model,
+                    google_api_key=google_api_key,
+                    groq_api_key=groq_api_key
+                )
                 
-            refs = "".join(item.get_ref().cref for item in items)
-            text = chunk.text
-            
-            document = Document(
-                page_content=text,
-                metadata={
-                    "doc_id": (doc_id := doc_id + 1),
-                    "source": source,
-                    "ref": refs,
-                }
-            )
-            texts.append(document)
-    
-    # Process tables (if any)
-    tables = []
-    for source, docling_document in conversions.items():
-        for table in docling_document.tables:
-            if table.label == DocItemLabel.TABLE:
-                ref = table.get_ref().cref
-                text = table.export_to_markdown()
+                st.session_state.embeddings_model = embeddings_model
+                st.session_state.embeddings_tokenizer = embeddings_tokenizer
+                st.session_state.vision_model = vision_model
+                st.session_state.llm_model = llm_model
+                st.session_state.models_loaded = True
                 
-                document = Document(
-                    page_content=text,
-                    metadata={
-                        "doc_id": (doc_id := doc_id + 1),
-                        "source": source,
-                        "ref": ref,
-                    },
-                )
-                tables.append(document)
-    
-    # Process images (if any)
-    pictures = []
-    start_doc_id = len(texts) + len(tables) + 1
-    
-    for source, docling_document in conversions.items():
-        if hasattr(docling_document, 'pictures') and docling_document.pictures:
-            for picture in docling_document.pictures:
-                try:
-                    ref = picture.get_ref().cref
-                    image = picture.get_image(docling_document)
-                    
-                    if image:
-                        response = vision_model.generate_content([
-                            "Extract all text and describe key visual elements in this image. "
-                            "Include any numbers, labels, or important details.",
-                            image
-                        ])
-                        
-                        document = Document(
-                            page_content=response.text,
-                            metadata={
-                                "doc_id": doc_id,
-                                "source": source,
-                                "ref": ref,
-                            }
-                        )
-                        pictures.append(document)
-                        doc_id += 1
-                except Exception as e:
-                    print(f"Error processing image: {str(e)}")
-    
-    return texts + tables + pictures
+                st.success("Models loaded successfully!")
+            except Exception as e:
+                st.error(f"Error loading models: {str(e)}")
 
-def create_vector_store(docs, embeddings_model):
-    # Create vector store (using Milvus as in your notebook)
-    # Note: You'll need to have Milvus running
-    vector_store = Milvus.from_documents(
-        docs,
-        embeddings_model,
-        connection_args={"host": "127.0.0.1", "port": "19530"},
-        collection_name="pdf_manual"
-    )
-    return vector_store
+# Main app interface
+st.title("PDF RAG Pipeline")
+st.write("Upload a PDF and ask questions about its content")
 
-def main():
-    st.title("PDF Manual Chatbot")
-    
-    # Initialize components
-    embeddings_model, embeddings_tokenizer, model, vision_model = initialize_components()
-    
-    # File upload
-    uploaded_file = st.file_uploader("Upload a PDF manual", type="pdf")
-    
-    if uploaded_file is not None:
-        # Save the uploaded file
-        file_path = os.path.join("temp", uploaded_file.name)
-        os.makedirs("temp", exist_ok=True)
-        with open(file_path, "wb") as f:
-            f.write(uploaded_file.getbuffer())
-        
-        # Process the PDF
-        with st.spinner("Processing PDF..."):
-            docs = process_pdf(file_path, embeddings_tokenizer, vision_model)
-            vector_store = create_vector_store(docs, embeddings_model)
-        
-        st.success("PDF processed successfully!")
-        
-        # Initialize chat history
-        if "messages" not in st.session_state:
-            st.session_state.messages = []
-        
-        # Display chat messages from history on app rerun
-        for message in st.session_state.messages:
-            with st.chat_message(message["role"]):
-                st.markdown(message["content"])
-        
-        # Accept user input
-        if prompt := st.chat_input("Ask a question about the manual"):
-            # Add user message to chat history
-            st.session_state.messages.append({"role": "user", "content": prompt})
-            
-            # Display user message in chat message container
-            with st.chat_message("user"):
-                st.markdown(prompt)
+if uploaded_file and st.session_state.models_loaded:
+    with st.spinner("Processing PDF..."):
+        try:
+            # Save uploaded file temporarily
+            file_path = f"./temp_{uploaded_file.name}"
+            with open(file_path, "wb") as f:
+                f.write(uploaded_file.getbuffer())
             
-            # Create QA chain
-            qa_chain = RetrievalQA.from_chain_type(
-                llm=model,
-                chain_type="stuff",
-                retriever=vector_store.as_retriever(),
-                return_source_documents=True
+            # Process the PDF
+            texts, tables, pictures = process_pdf(
+                file_path,
+                st.session_state.embeddings_tokenizer,
+                st.session_state.vision_model
             )
             
-            # Get response
-            with st.spinner("Thinking..."):
-                result = qa_chain({"query": prompt})
-                response = result["result"]
-                source_docs = result["source_documents"]
+            st.session_state.processed_docs = {
+                "texts": texts,
+                "tables": tables,
+                "pictures": pictures
+            }
             
-            # Display assistant response in chat message container
-            with st.chat_message("assistant"):
-                st.markdown(response)
-                
-                # Show sources if available
-                if source_docs:
-                    with st.expander("Source Documents"):
-                        for i, doc in enumerate(source_docs):
-                            st.write(f"Source {i+1}:")
-                            st.write(doc.page_content)
-                            st.write(f"Metadata: {doc.metadata}")
-                            st.write("---")
+            st.success("PDF processed successfully!")
+            
+            # Display document stats
+            col1, col2, col3 = st.columns(3)
+            col1.metric("Text Chunks", len(texts))
+            col2.metric("Tables", len(tables))
+            col3.metric("Images", len(pictures))
             
-            # Add assistant response to chat history
-            st.session_state.messages.append({"role": "assistant", "content": response})
+            # Remove temp file
+            os.remove(file_path)
+            
+        except Exception as e:
+            st.error(f"Error processing PDF: {str(e)}")
 
-if __name__ == "__main__":
-    main()
+# Question answering section
+if st.session_state.processed_docs:
+    st.divider()
+    st.subheader("Ask a Question")
+    
+    question = st.text_input("Enter your question about the document:")
+    
+    if question and st.button("Get Answer"):
+        with st.spinner("Generating answer..."):
+            try:
+                answer = query_pipeline(
+                    question=question,
+                    texts=st.session_state.processed_docs["texts"],
+                    tables=st.session_state.processed_docs["tables"],
+                    pictures=st.session_state.processed_docs["pictures"],
+                    embeddings_model=st.session_state.embeddings_model,
+                    llm_model=st.session_state.llm_model
+                )
+                
+                st.subheader("Answer")
+                st.write(answer)
+                
+            except Exception as e:
+                st.error(f"Error generating answer: {str(e)}")