Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,82 +1,64 @@
-# app_hf_space.py (Iteration)
 import streamlit as st
+import zipfile, io, os
+from langchain.document_loaders import PyPDFLoader
+from langchain.text_splitter import CharacterTextSplitter
+from langchain.schema import Document
+from langchain.embeddings import SentenceTransformerEmbeddings
+from langchain.vectorstores import FAISS
+from transformers import pipeline
 
-from langchain_huggingface import HuggingFaceEmbeddingsceEmbeddings
-from langchain_community.vectorstores.faiss import FAISS
-from langchain.chains import RetrievalQA
-from langchain_huggingface.llms import HuggingFacePipeline
+# Cache the QA initialization so ingestion runs once per session
+@st.cache_resource
+def init_qa(zip_bytes):
+    tmp_dir = "tmp_pdfs"
+    # Clean up or create temp folder
+    if os.path.exists(tmp_dir):
+        for f in os.listdir(tmp_dir):
+            os.remove(os.path.join(tmp_dir, f))
+    else:
+        os.makedirs(tmp_dir)
 
-# --- Configuration ---
-# Use Hugging Face inference API via the `pipeline` from langchain-huggingface
-EMBED_MODEL_ID = "sentence-transformers/all-MiniLM-L6-v2"
-LLM_MODEL_ID = "google/flan-t5-small"
-INDEX_DIR = "faiss_index"
+    # Extract uploaded ZIP
+    with zipfile.ZipFile(io.BytesIO(zip_bytes)) as z:
+        z.extractall(tmp_dir)
 
-# 1. Initialize embeddings via HF Inference API
-embeddings = HuggingFaceEmbeddings(
-    model_name=EMBED_MODEL_ID,
-    cache_dir=".hf_cache"
-)
+    # Load all PDFs
+    docs = []
+    for fname in os.listdir(tmp_dir):
+        if fname.lower().endswith(".pdf"):
+            loader = PyPDFLoader(os.path.join(tmp_dir, fname))
+            docs.extend(loader.load())
 
-# 2. Preload & Ingest (optional)
-import os, zipfile
-from streamlit import sidebar
+    # Split into manageable chunks
+    splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    split_docs = splitter.split_documents(docs)
 
-# Auto-extract preloaded zip if present
-docs_dir = "docs"
-zip_path = "preloaded_docs.zip"
-if os.path.exists(zip_path):
-    with zipfile.ZipFile(zip_path, "r") as z:
-        z.extractall(docs_dir)
-    sidebar.success(f"Extracted {zip_path} to {docs_dir}/")
+    # Build vector store
+    embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
+    vector_store = FAISS.from_documents(split_docs, embeddings)
 
-# Sidebar button to re-ingest docs and rebuild index
-if sidebar.button("Re-ingest docs & rebuild index"):
-    from ingest import load_documents, text_splitter, embeddings as ingest_embeddings
-    docs = load_documents(docs_dir)
-    chunks = text_splitter.split_documents(docs)
-    FAISS.from_documents(chunks, ingest_embeddings).save_local(INDEX_DIR)
-    sidebar.success("Re-ingestion complete and index rebuilt.")
+    # Load the RAG model
+    generator = pipeline(
+        "text2text-generation",
+        model="PleIAs/Pleias-RAG-350M",
+        tokenizer="PleIAs/Pleias-RAG-350M"
+    )
 
-# 3. Load FAISS index
-store = FAISS.load_local(
-    INDEX_DIR,
-    embeddings
-)
+    return vector_store, generator
 
-# 4. Initialize HF LLM via pipeline (inference API)
-# ...
-# 5. Build RetrievalQA chain
-# ...
-# 6. Streamlit UI
-# ... Initialize HF LLM via pipeline (inference API)
-llm = HuggingFacePipeline.from_model_id(
-    model_id=LLM_MODEL_ID,
-    task="text2text-generation",
-    pipeline_kwargs={
-        # Device mapping for inference
-        "device": -1,
-        # Cache directory for model weights
-        "cache_dir": ".hf_cache"
-    }
-)
-
-# 4. Build RetrievalQA chain
-aqa_chain = RetrievalQA.from_chain_type(
-    llm=llm,
-    chain_type="stuff",
-    retriever=store.as_retriever()
-)
-
-# 5. Streamlit UI
-def main():
-    st.title("🦜🔗 RAG App via HF Spaces")
-    query = st.text_input("Ask a question about your docs:")
+# Streamlit UI
+st.title("Pleias-RAG 350M Streamlit App")
+st.write("Upload a ZIP of PDFs to initialize the RAG engine.")
+zip_file = st.file_uploader("ZIP file", type=["zip"])
 
+if zip_file:
+    vector_store, generator = init_qa(zip_file.read())
+    query = st.text_input("Ask a question:")
     if query:
-        with st.spinner("Generating answer via HF Space..."):
-            answer = aqa_chain.run(query)
-        st.markdown(f"**Answer:** {answer}")
-
-if __name__ == "__main__":
-    main()
+        docs = vector_store.similarity_search(query, k=4)
+        context = "\n\n".join([doc.page_content for doc in docs])
+        prompt = f"question: {query}\ncontext: {context}"
+        answer = generator(prompt, max_length=512, do_sample=False)[0]["generated_text"]
+        st.write(answer)
+else:
+    st.info("Awaiting ZIP upload.")