Spaces:
Runtime error
Runtime error
File size: 3,822 Bytes
1fb1192 0f0a517 1fb1192 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | import os
from pathlib import Path
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
# Path to local PDFs
PDF_FILES = [
"1_agro.pdf",
"2-agro.pdf",
"3-agrogpt.pdf",
"4-agrogpt.pdf"
]
# Writeable path on Hugging Face for the vector store
VECTORSTORE_DIR = "vectorstore"
# Global variable to hold the FAISS index in memory
_vector_store = None
def get_embeddings_model():
"""Return the HuggingFace embeddings model. Uses a lightweight fast model."""
# Using all-MiniLM-L6-v2 as it's very fast and effective for semantic search
return HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
def initialize_knowledge_base():
"""
Checks if the vector database exists in /tmp/vectorstore.
If not, it reads the PDFs, chunks them, generates embeddings, and saves the DB.
"""
global _vector_store
if os.path.exists(VECTORSTORE_DIR) and os.path.exists(os.path.join(VECTORSTORE_DIR, "index.faiss")):
print(f"Loading existing vector database from {VECTORSTORE_DIR}...")
_vector_store = FAISS.load_local(VECTORSTORE_DIR, get_embeddings_model(), allow_dangerous_deserialization=True)
return _vector_store
print("Vector database not found. Initializing knowledge base (this may take a minute)...", flush=True)
os.makedirs(VECTORSTORE_DIR, exist_ok=True)
documents = []
for pdf_file in PDF_FILES:
try:
if os.path.exists(pdf_file):
print(f"Parsing {pdf_file}...", flush=True)
loader = PyPDFLoader(pdf_file)
documents.extend(loader.load())
else:
print(f"Warning: {pdf_file} not found in the root directory.", flush=True)
except Exception as e:
print(f"Error parsing {pdf_file}: {e}", flush=True)
if not documents:
print("No documents were loaded. Vector database initialization skipped.", flush=True)
return None
print(f"Total pages loaded: {len(documents)}. Splitting text...", flush=True)
# Split the documents into manageable chunks
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=1000,
chunk_overlap=200,
length_function=len
)
chunks = text_splitter.split_documents(documents)
print(f"Created {len(chunks)} text chunks. Generating embeddings...", flush=True)
# Generate embeddings and build FAISS index
embeddings = get_embeddings_model()
_vector_store = FAISS.from_documents(chunks, embeddings)
# Save for subsequent requests
_vector_store.save_local(VECTORSTORE_DIR)
print(f"Vector database built and saved to {VECTORSTORE_DIR} successfully.", flush=True)
return _vector_store
def query_rag(question: str, k: int = 3) -> str:
"""
Searches the FAISS vector database for chunks related to the question.
Returns a formatted string containing the retrieved context.
"""
global _vector_store
if not _vector_store:
_vector_store = initialize_knowledge_base()
if not _vector_store:
# If it's still None (e.g. PDFs missing or error), return empty context
return ""
try:
# Perform similarity search
docs = _vector_store.similarity_search(question, k=k)
# Format the retrieved documents into a single context string
context = "\n\n".join([f"[Source: {doc.metadata.get('source', 'Unknown')}]\n{doc.page_content}" for doc in docs])
return context
except Exception as e:
print(f"Error querying RAG: {e}", flush=True)
return ""
|