AsyncRAG

Sleeping

App Files Files Community

Zubaish commited on 8 days ago

Commit

2fd2129

1 Parent(s): b6d77d3

Final stable HF RAG (dataset-backed, CPU-safe)

Browse files

Files changed (5) hide show

app.py +11 -5
config.py +2 -13
frontend/index.html +19 -5
rag.py +50 -79
requirements.txt +14 -6

app.py CHANGED Viewed

@@ -1,3 +1,5 @@
 from fastapi import FastAPI
 from fastapi.responses import HTMLResponse
 from pydantic import BaseModel
@@ -5,17 +7,21 @@ from rag import ask_rag_with_status
 app = FastAPI()
 class Query(BaseModel):
     question: str
-# Serve frontend
 @app.get("/", response_class=HTMLResponse)
 def index():
-    with open("index.html", "r", encoding="utf-8") as f:
         return f.read()
-# Chat endpoint
 @app.post("/chat")
 def chat(q: Query):
-    result = ask_rag_with_status(q.question)
-    return result

+# app.py
 from fastapi import FastAPI
 from fastapi.responses import HTMLResponse
 from pydantic import BaseModel
 app = FastAPI()
 class Query(BaseModel):
     question: str
 @app.get("/", response_class=HTMLResponse)
 def index():
+    with open("frontend/index.html", "r", encoding="utf-8") as f:
         return f.read()
 @app.post("/chat")
 def chat(q: Query):
+    answer, status = ask_rag_with_status(q.question)
+    return {
+        "answer": answer,
+        "status": status,
+    }

config.py CHANGED Viewed

@@ -1,16 +1,5 @@
-import os
-# Hugging Face dataset repo containing PDFs
 HF_DATASET_REPO = "Zubaish/hubrag-kb"
-# Embedding model (lightweight, CPU-safe)
 EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
-# Chroma persistence (local to container)
-CHROMA_DIR = "/tmp/chroma"
-# LLM via HF Inference API (NOT local)
-LLM_MODEL = "microsoft/Phi-3-mini-4k-instruct"
-# Safety
-MAX_CONTEXT_CHUNKS = 4

+# config.py
 HF_DATASET_REPO = "Zubaish/hubrag-kb"
 EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
+LLM_MODEL = "google/flan-t5-small"  # SAFE on HF CPU

frontend/index.html CHANGED Viewed

@@ -4,17 +4,31 @@
   <meta charset="UTF-8" />
   <title>HubRAG</title>
   <style>
-    body { font-family: sans-serif; max-width: 800px; margin: 40px auto; }
-    textarea { width: 100%; padding: 10px; }
-    button { margin-top: 10px; padding: 8px 16px; }
-    pre { background: #f5f5f5; padding: 10px; white-space: pre-wrap; }
   </style>
 </head>
 <body>
 <h2>📄 HubRAG (HF Space)</h2>
-<textarea id="q" rows="4" placeholder="Ask a question..."></textarea>
 <br/>
 <button onclick="ask()">Ask</button>

   <meta charset="UTF-8" />
   <title>HubRAG</title>
   <style>
+    body {
+      font-family: sans-serif;
+      max-width: 800px;
+      margin: 40px auto;
+    }
+    textarea {
+      width: 100%;
+      padding: 10px;
+    }
+    button {
+      margin-top: 10px;
+      padding: 8px 16px;
+    }
+    pre {
+      background: #f5f5f5;
+      padding: 10px;
+      white-space: pre-wrap;
+    }
   </style>
 </head>
 <body>
 <h2>📄 HubRAG (HF Space)</h2>
+<textarea id="q" rows="4" placeholder="Ask a question about the documents..."></textarea>
 <br/>
 <button onclick="ask()">Ask</button>

rag.py CHANGED Viewed

@@ -1,99 +1,78 @@
-import os
 from datasets import load_dataset
-from langchain_community.document_loaders import PyPDFLoader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import Chroma
-from langchain.embeddings import HuggingFaceEmbeddings
-from huggingface_hub import InferenceClient
-from config import (
-    HF_DATASET_REPO,
-    EMBEDDING_MODEL,
-    CHROMA_DIR,
-    LLM_MODEL,
-    MAX_CONTEXT_CHUNKS,
-)
-# --- Globals (lazy loaded) ---
-_vectordb = None
-# --- Embeddings ---
-embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)
-# --- HF Inference Client ---
-llm = InferenceClient(
-    model=LLM_MODEL,
-    token=os.environ.get("HF_TOKEN"),
-)
-# --- Load PDFs from HF Dataset ---
 def load_documents():
-    docs = []
     ds = load_dataset(HF_DATASET_REPO, split="train")
-    for item in ds:
-        pdf_path = item["file"]
-        loader = PyPDFLoader(pdf_path)
-        docs.extend(loader.load())
     return docs
-def get_vectordb():
-    global _vectordb
-    if _vectordb is not None:
-        return _vectordb
-    documents = load_documents()
-    if not documents:
-        return None
-    splitter = RecursiveCharacterTextSplitter(
-        chunk_size=800,
-        chunk_overlap=150
-    )
-    chunks = splitter.split_documents(documents)
-    if not chunks:
-        return None
-    _vectordb = Chroma.from_documents(
-        chunks,
-        embedding=embeddings,
-        persist_directory=CHROMA_DIR
-    )
-    return _vectordb
-# --- RAG Query ---
 def ask_rag_with_status(question: str):
     status = []
-    vectordb = get_vectordb()
-    if vectordb is None:
-        return {
-            "answer": "No documents indexed.",
-            "status": ["Vector DB not available"]
-        }
-    status.append("🔍 Searching documents")
-    docs = vectordb.similarity_search(question, k=MAX_CONTEXT_CHUNKS)
     if not docs:
-        return {
-            "answer": "No relevant context found.",
-            "status": status
-        }
     context = "\n\n".join(d.page_content for d in docs)
-    prompt = f"""You are a helpful assistant.
-Answer ONLY from the context below.
-If the answer is not present, say "I don't know".
 Context:
 {context}
@@ -105,14 +84,6 @@ Answer:
 """
     status.append("🧠 Generating answer")
-    answer = llm.text_generation(
-        prompt,
-        max_new_tokens=256,
-        temperature=0.2,
-    )
-    return {
-        "answer": answer.strip(),
-        "status": status
-    }

+# rag.py
 from datasets import load_dataset
 from langchain_community.vectorstores import Chroma
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain.schema import Document
+from transformers import pipeline
+from config import HF_DATASET_REPO, EMBEDDING_MODEL, LLM_MODEL
+# ------------------------
+# Load documents from HF Dataset
+# ------------------------
 def load_documents():
     ds = load_dataset(HF_DATASET_REPO, split="train")
+    docs = []
+    for row in ds:
+        text = row.get("text") or row.get("content")
+        if text and text.strip():
+            docs.append(Document(page_content=text))
     return docs
+# ------------------------
+# Build Vector DB (ONCE)
+# ------------------------
+documents = load_documents()
+if not documents:
+    raise RuntimeError("No documents loaded from HF Dataset")
+splitter = RecursiveCharacterTextSplitter(
+    chunk_size=500,
+    chunk_overlap=50,
+)
+chunks = splitter.split_documents(documents)
+embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)
+vectordb = Chroma.from_documents(
+    documents=chunks,
+    embedding=embeddings,
+)
+retriever = vectordb.as_retriever(search_kwargs={"k": 3})
+# ------------------------
+# LLM (CPU SAFE)
+# ------------------------
+llm = pipeline(
+    "text2text-generation",
+    model=LLM_MODEL,
+    max_new_tokens=256,
+)
+# ------------------------
+# RAG Query
+# ------------------------
 def ask_rag_with_status(question: str):
     status = []
+    status.append("🔎 Retrieving documents")
+    docs = retriever.get_relevant_documents(question)
     if not docs:
+        return "No relevant documents found.", status
     context = "\n\n".join(d.page_content for d in docs)
+    prompt = f"""
+Answer the question using the context below.
 Context:
 {context}
 """
     status.append("🧠 Generating answer")
+    result = llm(prompt)[0]["generated_text"]
+    return result.strip(), status

requirements.txt CHANGED Viewed

@@ -1,10 +1,18 @@
 fastapi
 uvicorn
 pydantic
-datasets
-huggingface_hub
 sentence-transformers
-langchain
-langchain-community
-chromadb
-pypdf

 fastapi
 uvicorn
 pydantic
+python-dotenv
+langchain==0.2.17
+langchain-community==0.2.17
+langchain-text-splitters==0.2.4
+chromadb==0.5.5
 sentence-transformers
+pypdf
+transformers>=4.39.0
+huggingface_hub
+datasets
+torch