Update app.py

app.py

@@ -1,14 +1,17 @@
 import os
-import pickle
+import streamlit as st
 import PyPDF2
 import numpy as np
 import faiss
-import torch
-import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM, DataCollatorForLanguageModeling, Trainer, TrainingArguments
-from datasets import Dataset
 from sentence_transformers import SentenceTransformer
-from peft import LoraConfig, get_peft_model
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from langchain.chains import ConversationalRetrievalChain
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import FAISS
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.llms import HuggingFacePipeline
+from langchain.prompts import PromptTemplate
+from transformers import pipeline
 
 # Load embedding model
 @st.cache_resource
@@ -24,47 +27,15 @@ def parse_pdf(file):
     return text
 
 # Split text into chunks
-def split_text(text, chunk_size=500):
-    return [text[i:i + chunk_size] for i in range(0, len(text), chunk_size)]
+def split_text(text, chunk_size=500, chunk_overlap=100):
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
+    return text_splitter.split_text(text)
 
 # Create FAISS index
-def create_faiss_index(embeddings):
-    dimension = embeddings.shape[1]
-    index = faiss.IndexFlatL2(dimension)
-    index.add(embeddings)
-    return index
-
-# Fine-tune the model
-def fine_tune_model(dataset, model_name, output_dir="./fine-tuned-model"):
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name)
-
-    def preprocess_function(examples):
-        inputs = [f"Question: {q} Answer: {a}" for q, a in zip(examples["question"], examples["answer"])]
-        return tokenizer(inputs, truncation=True, padding="max_length", max_length=512)
-
-    tokenized_dataset = dataset.map(preprocess_function, batched=True)
-
-    training_args = TrainingArguments(
-        output_dir=output_dir,
-        per_device_train_batch_size=4,
-        num_train_epochs=3,
-        save_steps=10_000,
-        save_total_limit=2,
-    )
-
-    trainer = Trainer(
-        model=model,
-        args=training_args,
-        train_dataset=tokenized_dataset,
-        tokenizer=tokenizer,
-        data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
-    )
-
-    trainer.train()
-    model.save_pretrained(output_dir)
-    tokenizer.save_pretrained(output_dir)
-    return output_dir
+def create_faiss_index(texts, embedding_model):
+    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    vectorstore = FAISS.from_texts(texts, embeddings)
+    return vectorstore
 
 # Generate response from the model
 def generate_response(prompt, model, tokenizer):
@@ -75,62 +46,89 @@ def generate_response(prompt, model, tokenizer):
 
 # Main Streamlit app
 def main():
-    st.title("Chat with PDF using Fine-Tuned Llama Model")
-
-    # Step 1: Upload PDF file
-    uploaded_file = st.file_uploader("Upload a PDF file", type=["pdf"])
-    if uploaded_file is not None:
-        st.write("File uploaded successfully!")
-
-        # Process PDF
-        with st.spinner("Processing PDF..."):
-            pdf_text = parse_pdf(uploaded_file)
-            chunks = split_text(pdf_text)
+    st.title("Advanced Chat with Your Document")
+
+    # Initialize session state for conversation history and documents
+    if "conversation_history" not in st.session_state:
+        st.session_state.conversation_history = []
+    if "vectorstore" not in st.session_state:
+        st.session_state.vectorstore = None
+
+    # Step 1: Upload multiple PDF files
+    uploaded_files = st.file_uploader("Upload PDF files", type=["pdf"], accept_multiple_files=True)
+    if uploaded_files:
+        st.write(f"{len(uploaded_files)} file(s) uploaded successfully!")
+
+        # Process PDFs
+        with st.spinner("Processing PDFs..."):
+            all_texts = []
+            for uploaded_file in uploaded_files:
+                pdf_text = parse_pdf(uploaded_file)
+                chunks = split_text(pdf_text)
+                all_texts.extend(chunks)
+
+            # Create a unified vector database
             embedding_model = load_embedding_model()
-            chunk_embeddings = embedding_model.encode(chunks)
-            faiss_index = create_faiss_index(np.array(chunk_embeddings))
-
-        st.success("PDF processed! Proceed to fine-tuning.")
-
-        # Step 2: Fine-tune the model
-        if st.button("Fine-Tune Model"):
-            with st.spinner("Fine-tuning the model..."):
-                # Create a dataset of question-answer pairs
-                qa_pairs = []
-                for chunk in chunks:
-                    qa_pairs.append({"question": "What is this about?", "answer": chunk[:100]})  # Simplified example
-
-                dataset = Dataset.from_dict({
-                    "question": [pair["question"] for pair in qa_pairs],
-                    "answer": [pair["answer"] for pair in qa_pairs],
-                })
-
-                # Fine-tune the model
-                model_name = "meta-llama/Llama-2-7b-chat-hf"  # Replace with your local path
-                fine_tuned_model_path = fine_tune_model(dataset, model_name)
-
-            st.success(f"Model fine-tuned! Saved at: {fine_tuned_model_path}")
-
-            # Load the fine-tuned model
-            tokenizer = AutoTokenizer.from_pretrained(fine_tuned_model_path)
-            model = AutoModelForCausalLM.from_pretrained(fine_tuned_model_path, device_map="auto", torch_dtype=torch.float16)
-
-            st.success("Fine-tuned model loaded! You can now ask questions.")
-
-        # Step 3: Chat interface
-        user_input = st.text_input("Ask a question about the PDF:")
-        if user_input:
-            with st.spinner("Generating response..."):
-                # Retrieve relevant chunk
-                query_embedding = embedding_model.encode([user_input])
-                _, indices = faiss_index.search(query_embedding, k=1)
-                relevant_chunk = chunks[indices[0][0]]
-
-                # Generate response
-                prompt = f"Context: {relevant_chunk}\nQuestion: {user_input}\nAnswer:"
-                response = generate_response(prompt, model, tokenizer)
-
-            st.write(f"**Response:** {response}")
+            st.session_state.vectorstore = create_faiss_index(all_texts, embedding_model)
+
+        st.success("PDFs processed! You can now ask questions.")
+
+    # Step 2: Chat interface
+    user_input = st.text_input("Ask a question about the document(s):")
+    if user_input:
+        if st.session_state.vectorstore is None:
+            st.error("Please upload and process documents first.")
+            return
+
+        with st.spinner("Generating response..."):
+            # Load the LLM
+            model_name = "meta-llama/Llama-2-7b-chat-hf"  # Replace with your local path
+            tokenizer = AutoTokenizer.from_pretrained(model_name)
+            model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16)
+
+            # Set up LangChain components
+            retriever = st.session_state.vectorstore.as_retriever()
+            llm = HuggingFacePipeline(pipeline=pipeline("text-generation", model=model, tokenizer=tokenizer))
+
+            # Define a custom prompt template for Chain-of-Thought reasoning
+            prompt_template = """
+            Answer the following question based ONLY on the provided context.
+            If the question requires multi-step reasoning, break it down step by step.
+            Context: {context}
+            Question: {question}
+            Answer:
+            """
+            prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
+
+            # Create a conversational retrieval chain
+            qa_chain = ConversationalRetrievalChain.from_llm(
+                llm=llm,
+                retriever=retriever,
+                combine_docs_chain_kwargs={"prompt": prompt},
+                return_source_documents=True
+            )
+
+            # Add conversation history
+            chat_history = st.session_state.conversation_history[-3:]  # Last 3 interactions
+            result = qa_chain({"question": user_input, "chat_history": chat_history})
+
+            # Extract response and update conversation history
+            response = result["answer"]
+            st.session_state.conversation_history.append(f"User: {user_input}")
+            st.session_state.conversation_history.append(f"Bot: {response}")
+
+        st.write(f"**Response:** {response}")
+
+        # Display source documents (optional)
+        if "source_documents" in result:
+            st.subheader("Source Documents")
+            for doc in result["source_documents"]:
+                st.write(doc.page_content)
+
+        # Display conversation history
+        st.subheader("Conversation History")
+        for line in st.session_state.conversation_history:
+            st.write(line)
 
 if __name__ == "__main__":
     main()