Update app.py

app.py

@@ -1,103 +1,127 @@
-import streamlit as st
-from langchain_community.document_loaders import PyPDFLoader, TextLoader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.vectorstores import FAISS
-from langchain.llms import HuggingFacePipeline
-from langchain.chains import RetrievalQA
-
-from transformers import pipeline
-
-# -------------------------------
-# Load Documents
-# -------------------------------
-def load_documents(uploaded_files):
-    documents = []
-    for file in uploaded_files:
-        with open(file.name, "wb") as f:
-            f.write(file.getbuffer())
-
-        if file.name.endswith(".pdf"):
-            loader = PyPDFLoader(file.name)
-        else:
-            loader = TextLoader(file.name)
-
-        documents.extend(loader.load())
-    return documents
-
-
-# -------------------------------
-# Split Documents
-# -------------------------------
-def split_documents(documents):
-    splitter = RecursiveCharacterTextSplitter(
-        chunk_size=500,
-        chunk_overlap=50
-    )
-    return splitter.split_documents(documents)
-
-
-# -------------------------------
-# Create Vector Store
-# -------------------------------
-def create_vectorstore(chunks):
-    embeddings = HuggingFaceEmbeddings(
-        model_name="sentence-transformers/all-MiniLM-L6-v2"
-    )
-    return FAISS.from_documents(chunks, embeddings)
-
-
-# -------------------------------
-# Load Local LLM (FREE)
-# -------------------------------
-def load_llm():
-    pipe = pipeline(
-        "text-generation",
-        model="google/flan-t5-base",
-        max_length=512
-    )
-    return HuggingFacePipeline(pipeline=pipe)
-
-
-# -------------------------------
-# Build QA Chain
-# -------------------------------
-def build_qa(vectorstore):
-    llm = load_llm()
-    retriever = vectorstore.as_retriever()
-
-    qa = RetrievalQA.from_chain_type(
-        llm=llm,
-        retriever=retriever,
-        return_source_documents=False
-    )
-    return qa
-
-
-# -------------------------------
-# Streamlit UI
-# -------------------------------
-st.set_page_config(page_title="RAG Chatbot", layout="wide")
-st.title("📄 Chat with Your Documents (RAG)")
-
-uploaded_files = st.file_uploader(
-    "Upload PDF or TXT files",
-    accept_multiple_files=True
-)
-
-if uploaded_files:
-    with st.spinner("Processing documents..."):
-        docs = load_documents(uploaded_files)
-        chunks = split_documents(docs)
-        vectorstore = create_vectorstore(chunks)
-        qa_chain = build_qa(vectorstore)
-
-    st.success("Documents ready!")
-
-    query = st.text_input("Ask a question from your documents")
-
-    if query:
-        with st.spinner("Generating answer..."):
-            result = qa_chain.run(query)
-            st.write("### Answer:")
-            st.write(result)
+import streamlit as st
+import tempfile
+
+from langchain_community.document_loaders import PyPDFLoader, TextLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import FAISS
+from langchain.llms import HuggingFacePipeline
+from langchain.chains import RetrievalQA
+
+from transformers import pipeline
+
+# -------------------------------
+# Page Config
+# -------------------------------
+st.set_page_config(page_title="RAG Chatbot", layout="wide")
+st.title("📄 Chat with Your Documents (RAG)")
+st.write("🚀 App started successfully")
+
+# -------------------------------
+# Load Documents (FIXED)
+# -------------------------------
+def load_documents(uploaded_files):
+    documents = []
+
+    for file in uploaded_files:
+        # Save file safely using temp file
+        with tempfile.NamedTemporaryFile(delete=False, suffix=file.name) as tmp:
+            tmp.write(file.getbuffer())
+            temp_path = tmp.name
+
+        # Load based on type
+        if file.name.endswith(".pdf"):
+            loader = PyPDFLoader(temp_path)
+        else:
+            loader = TextLoader(temp_path)
+
+        documents.extend(loader.load())
+
+    return documents
+
+
+# -------------------------------
+# Split Documents
+# -------------------------------
+def split_documents(documents):
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size=500,
+        chunk_overlap=50
+    )
+    return splitter.split_documents(documents)
+
+
+# -------------------------------
+# Cached Embeddings (IMPORTANT)
+# -------------------------------
+@st.cache_resource
+def get_embeddings():
+    return HuggingFaceEmbeddings(
+        model_name="sentence-transformers/all-MiniLM-L6-v2"
+    )
+
+
+# -------------------------------
+# Create Vector Store
+# -------------------------------
+def create_vectorstore(chunks):
+    embeddings = get_embeddings()
+    return FAISS.from_documents(chunks, embeddings)
+
+
+# -------------------------------
+# Cached LLM (IMPORTANT)
+# -------------------------------
+@st.cache_resource
+def load_llm():
+    pipe = pipeline(
+        "text-generation",
+        model="google/flan-t5-small",  # lightweight model
+        max_length=256
+    )
+    return HuggingFacePipeline(pipeline=pipe)
+
+
+# -------------------------------
+# Build QA Chain
+# -------------------------------
+def build_qa(vectorstore):
+    llm = load_llm()
+    retriever = vectorstore.as_retriever()
+
+    qa = RetrievalQA.from_chain_type(
+        llm=llm,
+        retriever=retriever,
+        return_source_documents=False
+    )
+    return qa
+
+
+# -------------------------------
+# UI - Upload
+# -------------------------------
+uploaded_files = st.file_uploader(
+    "Upload PDF or TXT files",
+    accept_multiple_files=True
+)
+
+if uploaded_files:
+    with st.spinner("📄 Processing documents..."):
+        docs = load_documents(uploaded_files)
+        chunks = split_documents(docs)
+        vectorstore = create_vectorstore(chunks)
+        qa_chain = build_qa(vectorstore)
+
+    st.success("✅ Documents ready!")
+
+    # -------------------------------
+    # User Query
+    # -------------------------------
+    query = st.text_input("💬 Ask a question from your documents")
+
+    if query:
+        with st.spinner("🤖 Generating answer..."):
+            result = qa_chain.run(query)
+
+        st.markdown("### 🧠 Answer:")
+        st.write(result)