update

ingest.py

@@ -1,42 +1,62 @@
-import os, shutil
-from datasets import load_dataset
+# ingest.py
+import os
+import shutil
+from huggingface_hub import hf_hub_download, list_repo_files
 from langchain_community.document_loaders import Docx2txtLoader
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_chroma import Chroma
-from config import KB_DIR, HF_DATASET_REPO, EMBEDDING_MODEL, CHROMA_DIR, CHUNK_SIZE, CHUNK_OVERLAP
+from config import KB_DIR, HF_DATASET_REPO, EMBEDDING_MODEL, CHROMA_DIR, CHUNK_SIZE, CHUNK_OVERLAP, HF_TOKEN
 
 def run_ingestion():
+    # 1. Clean Environment
     if os.path.exists(KB_DIR): shutil.rmtree(KB_DIR)
     if os.path.exists(CHROMA_DIR): shutil.rmtree(CHROMA_DIR)
     os.makedirs(KB_DIR, exist_ok=True)
 
-    print(f"⬇️ Loading dataset from {HF_DATASET_REPO}...")
-    dataset = load_dataset(HF_DATASET_REPO, split="train", decode=False)
+    print(f"⬇️ Downloading files from NEW repo: {HF_DATASET_REPO}...")
     
-    docs = []
-    for i, row in enumerate(dataset):
-        file_item = row.get("docx") or row.get("file")
-        src_path = file_item.get("path") if isinstance(file_item, dict) else None
+    try:
+        # List files using the hub API instead of load_dataset
+        all_files = list_repo_files(repo_id=HF_DATASET_REPO, repo_type="dataset", token=HF_TOKEN)
+        docx_files = [f for f in all_files if f.lower().endswith(".docx")]
+        
+        docs = []
+        for i, file_name in enumerate(docx_files):
+            # Download file directly to local folder
+            local_path = hf_hub_download(
+                repo_id=HF_DATASET_REPO,
+                filename=file_name,
+                repo_type="dataset",
+                local_dir=KB_DIR,
+                token=HF_TOKEN
+            )
+            
+            # Load the text from docx
+            loader = Docx2txtLoader(local_path)
+            docs.extend(loader.load())
+            print(f"✅ Loaded: {file_name}")
 
-        if src_path and src_path.lower().endswith(".docx"):
-            dest_path = os.path.join(KB_DIR, f"doc_{i}.docx")
-            shutil.copy(src_path, dest_path)
-            try:
-                loader = Docx2txtLoader(dest_path)
-                docs.extend(loader.load())
-                print(f"✅ Loaded: doc_{i}.docx")
-            except Exception as e:
-                print(f"❌ Error loading doc_{i}: {e}")
+        if not docs:
+            print("❌ No documents found. Check repo files.")
+            return
 
-    if not docs:
-        print("❌ CRITICAL: No .docx documents found.")
-        return
+        # 2. Chunking
+        splitter = RecursiveCharacterTextSplitter(chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP)
+        splits = splitter.split_documents(docs)
 
-    splits = RecursiveCharacterTextSplitter(chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP).split_documents(docs)
-    embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)
-    Chroma.from_documents(documents=splits, embedding=embeddings, persist_directory=CHROMA_DIR)
-    print(f"✅ Knowledge base initialized at {CHROMA_DIR}")
+        # 3. Embedding and Storage
+        print(f"🧠 Indexing {len(splits)} chunks into ChromaDB...")
+        embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)
+        Chroma.from_documents(
+            documents=splits, 
+            embedding=embeddings, 
+            persist_directory=CHROMA_DIR
+        )
+        print(f"✅ Knowledge base initialized successfully at {CHROMA_DIR}")
+
+    except Exception as e:
+        print(f"❌ Ingestion failed: {e}")
 
 if __name__ == "__main__":
     run_ingestion()

rag.py

@@ -1,26 +1,34 @@
+# rag.py
 import os
 from transformers import pipeline
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_chroma import Chroma
-from config import EMBEDDING_MODEL, LLM_MODEL, CHROMA_DIR, LLM_TASK
+from config import EMBEDDING_MODEL, LLM_MODEL, CHROMA_DIR
 
 embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)
 
-if os.path.exists(CHROMA_DIR) and os.path.isdir(CHROMA_DIR):
+# Check if directory exists AND has files
+if os.path.exists(CHROMA_DIR) and any(os.scandir(CHROMA_DIR)):
     vectordb = Chroma(persist_directory=CHROMA_DIR, embedding_function=embeddings)
-    print("✅ Vector DB loaded")
+    print("✅ Vector DB ready")
 else:
     vectordb = None
+    print("⚠️ Vector DB not found or empty")
 
-qa_pipeline = pipeline(task=LLM_TASK, model=LLM_MODEL, max_new_tokens=256, trust_remote_code=True)
+qa_pipeline = pipeline(
+    task="text-generation", 
+    model=LLM_MODEL, 
+    max_new_tokens=256,
+    trust_remote_code=True
+)
 
 def ask_rag_with_status(question: str):
     if vectordb is None:
-        return "Knowledge base not initialized. Check build logs.", "ERROR"
+        return "The knowledge base is not initialized. Please check deployment logs.", "ERROR"
 
     docs = vectordb.similarity_search(question, k=3)
     context = "\n\n".join(d.page_content for d in docs)
-    prompt = f"Answer using the context.\nContext: {context}\nQuestion: {question}\nAnswer:"
+    prompt = f"Context: {context}\n\nQuestion: {question}\nAnswer:"
     
     result = qa_pipeline(prompt)
     answer = result[0]["generated_text"].split("Answer:")[-1].strip()