Stable HF-ready RAG using HF Datasets

.gitignore

@@ -1,18 +1,4 @@
-# Python
-__pycache__/
-*.pyc
-
-# Chroma DB
-chroma_db/
-data/
-
-# Environment
-.env
-
-# Frontend (DO NOT COMMIT)
-frontend/node_modules/
-frontend/dist/
-frontend/.vite/
-
 kb/
-vectordb/
+chroma_db/
+*.pdf
+__pycache__/

Dockerfile

@@ -9,8 +9,6 @@ RUN pip install --no-cache-dir -r requirements.txt
 
 COPY app.py rag.py ingest.py config.py ./
 
-RUN mkdir -p kb vectordb
-
 EXPOSE 7860
 
 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -1,18 +1,13 @@
 from fastapi import FastAPI
-from pydantic import BaseModel
 from rag import ask_rag_with_status
 
-app = FastAPI(title="RAG Knowledge Bot")
-
-class Query(BaseModel):
-    question: str
-
+app = FastAPI()
 
 @app.get("/")
 def health():
     return {"status": "ok"}
 
-
-@app.post("/chat")
-def chat(query: Query):
-    return ask_rag_with_status(query.question)
+@app.post("/ask")
+def ask(payload: dict):
+    question = payload.get("question", "")
+    return ask_rag_with_status(question)

config.py

@@ -1,9 +1,18 @@
 import os
+from huggingface_hub import snapshot_download
+from config import HF_DATASET_ID, KB_DIR
 
-BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+def download_kb():
+    os.makedirs(KB_DIR, exist_ok=True)
 
-KB_DIR = os.path.join(BASE_DIR, "kb")
-VECTOR_DB_DIR = os.path.join(BASE_DIR, "vectordb")
+    snapshot_download(
+        repo_id=HF_DATASET_ID,
+        repo_type="dataset",
+        local_dir=KB_DIR,
+        local_dir_use_symlinks=False
+    )
 
-EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
-LLM_MODEL = "microsoft/Phi-3-mini-4k-instruct"
+    print("✅ Knowledge base downloaded")
+
+if __name__ == "__main__":
+    download_kb()

ingest.py

@@ -1,30 +1,18 @@
 import os
-from langchain_community.document_loaders import DirectoryLoader, PyPDFLoader
-from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain_community.vectorstores import Chroma
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from config import KB_DIR, VECTOR_DB_DIR, EMBEDDING_MODEL
+from huggingface_hub import snapshot_download
+from config import HF_DATASET_ID, KB_DIR
 
-def ingest():
-    if not os.path.exists(KB_DIR) or not os.listdir(KB_DIR):
-        print("⚠️ No PDFs found in kb/. Skipping ingestion.")
-        return
+def download_kb():
+    os.makedirs(KB_DIR, exist_ok=True)
 
-    loader = DirectoryLoader(KB_DIR, glob="**/*.pdf", loader_cls=PyPDFLoader)
-    docs = loader.load()
-
-    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
-    chunks = splitter.split_documents(docs)
-
-    embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)
-
-    Chroma.from_documents(
-        chunks,
-        embeddings,
-        persist_directory=VECTOR_DB_DIR
+    snapshot_download(
+        repo_id=HF_DATASET_ID,
+        repo_type="dataset",
+        local_dir=KB_DIR,
+        local_dir_use_symlinks=False
     )
 
-    print("✅ Ingestion complete")
+    print("✅ Knowledge base downloaded")
 
 if __name__ == "__main__":
-    ingest()
+    download_kb()

rag.py

@@ -1,74 +1,66 @@
-from langchain_community.vectorstores import Chroma
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
-
-from config import VECTOR_DB_DIR, EMBEDDING_MODEL, LLM_MODEL
-
-_embeddings = None
-_db = None
-_tokenizer = None
-_model = None
-
-
-def get_vector_db():
-    global _embeddings, _db
+import os
+from typing import Dict
 
-    if _db is None:
-        _embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)
-        _db = Chroma(
-            persist_directory=VECTOR_DB_DIR,
-            embedding_function=_embeddings,
-        )
-    return _db
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import Chroma
 
+from ingest import download_kb
+from config import (
+    KB_DIR,
+    CHROMA_DIR,
+    EMBED_MODEL,
+    CHUNK_SIZE,
+    CHUNK_OVERLAP,
+)
 
-def get_llm():
-    global _tokenizer, _model
+# -------------------------
+# Startup: download + index
+# -------------------------
 
-    if _model is None:
-        _tokenizer = AutoTokenizer.from_pretrained(
-            LLM_MODEL, trust_remote_code=True
-        )
-        _model = AutoModelForCausalLM.from_pretrained(
-            LLM_MODEL,
-            trust_remote_code=True,
-            torch_dtype=torch.float32
-        )
-    return _tokenizer, _model
+print("⬇️ Downloading KB...")
+download_kb()
 
+print("📄 Loading documents...")
+documents = []
+for file in os.listdir(KB_DIR):
+    if file.endswith(".pdf"):
+        loader = PyPDFLoader(os.path.join(KB_DIR, file))
+        documents.extend(loader.load())
 
-def ask_rag_with_status(question: str):
-    status = []
+print(f"📚 Loaded {len(documents)} pages")
 
-    db = get_vector_db()
-    status.append("📚 Vector DB loaded")
+splitter = RecursiveCharacterTextSplitter(
+    chunk_size=CHUNK_SIZE,
+    chunk_overlap=CHUNK_OVERLAP,
+)
 
-    docs = db.similarity_search(question, k=3)
-    context = "\n\n".join(d.page_content for d in docs)
-    status.append("🔍 Retrieved relevant context")
+splits = splitter.split_documents(documents)
 
-    tokenizer, model = get_llm()
-    status.append("🤖 LLM loaded")
+embeddings = HuggingFaceEmbeddings(
+    model_name=EMBED_MODEL
+)
 
-    prompt = f"""
-You are a helpful assistant.
+vectorstore = Chroma.from_documents(
+    documents=splits,
+    embedding=embeddings,
+    persist_directory=CHROMA_DIR
+)
 
-Context:
-{context}
+retriever = vectorstore.as_retriever(search_kwargs={"k": 4})
 
-Question:
-{question}
+# -------------------------
+# Query API
+# -------------------------
 
-Answer clearly and concisely.
-"""
+def ask_rag_with_status(question: str) -> Dict:
+    docs = retriever.get_relevant_documents(question)
 
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_new_tokens=300)
-
-    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    context = "\n\n".join(d.page_content for d in docs)
 
     return {
-        "answer": answer,
-        "status": status
+        "question": question,
+        "chunks_used": len(docs),
+        "context_preview": context[:500]
     }

requirements.txt

@@ -1,17 +1,10 @@
 fastapi
 uvicorn
-python-dotenv
-
 langchain==0.2.17
 langchain-community==0.2.17
 langchain-text-splitters==0.2.4
-
 chromadb==0.5.5
 sentence-transformers
+huggingface_hub
 pypdf
-
-transformers>=4.39.0
-huggingface_hub<1.0.0
-numpy<2
-SQLAlchemy<3
-requests<3
+numpy<2