test RAG with local Ollama models
initial commit
This commit is contained in:
commit
b1405f3b84
25
.gitignore
vendored
Normal file
25
.gitignore
vendored
Normal file
@ -0,0 +1,25 @@
|
|||||||
|
# Python
|
||||||
|
__pycache__/
|
||||||
|
*.py[cod]
|
||||||
|
*$py.class
|
||||||
|
*.so
|
||||||
|
.Python
|
||||||
|
venv/
|
||||||
|
env/
|
||||||
|
ENV/
|
||||||
|
|
||||||
|
# Vector store
|
||||||
|
vectorstore/
|
||||||
|
|
||||||
|
# Uploads
|
||||||
|
uploads/
|
||||||
|
|
||||||
|
# IDE
|
||||||
|
.vscode/
|
||||||
|
.idea/
|
||||||
|
*.swp
|
||||||
|
*.swo
|
||||||
|
|
||||||
|
# OS
|
||||||
|
.DS_Store
|
||||||
|
Thumbs.db
|
||||||
65
README.md
Normal file
65
README.md
Normal file
@ -0,0 +1,65 @@
|
|||||||
|
# Local RAG Setup
|
||||||
|
|
||||||
|
Minimal RAG implementation with LangChain, Ollama, and FAISS.
|
||||||
|
|
||||||
|
## Dependencies
|
||||||
|
|
||||||
|
Only 5 packages:
|
||||||
|
- `langchain` - Core framework
|
||||||
|
- `langchain-ollama` - Ollama integration
|
||||||
|
- `faiss-cpu` - Vector search
|
||||||
|
- `sentence-transformers` - Embeddings
|
||||||
|
- `pypdf` - PDF loading
|
||||||
|
|
||||||
|
## Installation
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Create conda environment
|
||||||
|
conda create -n local_rag python=3.10 -y
|
||||||
|
conda activate local_rag
|
||||||
|
|
||||||
|
# Install dependencies
|
||||||
|
pip install -r requirements.txt
|
||||||
|
```
|
||||||
|
|
||||||
|
## Setup Ollama
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Make sure Ollama is running
|
||||||
|
ollama serve
|
||||||
|
|
||||||
|
# Pull a model (in another terminal)
|
||||||
|
ollama pull llama2
|
||||||
|
```
|
||||||
|
|
||||||
|
## Usage
|
||||||
|
|
||||||
|
Edit `local_rag.py` and uncomment the example code:
|
||||||
|
|
||||||
|
```python
|
||||||
|
# Add documents
|
||||||
|
rag.add_documents([
|
||||||
|
"path/to/document1.pdf",
|
||||||
|
"path/to/document2.txt"
|
||||||
|
])
|
||||||
|
|
||||||
|
# Query
|
||||||
|
question = "What is this document about?"
|
||||||
|
answer = rag.query(question)
|
||||||
|
print(f"Answer: {answer}")
|
||||||
|
```
|
||||||
|
|
||||||
|
Run:
|
||||||
|
```bash
|
||||||
|
python local_rag.py
|
||||||
|
```
|
||||||
|
|
||||||
|
## How it works
|
||||||
|
|
||||||
|
1. **Load documents** - PDFs or text files
|
||||||
|
2. **Split into chunks** - 1000 chars with 200 overlap
|
||||||
|
3. **Create embeddings** - Using sentence-transformers
|
||||||
|
4. **Store in FAISS** - Fast similarity search
|
||||||
|
5. **Query** - Retrieve relevant chunks and generate answer with Ollama
|
||||||
|
|
||||||
|
That's it! Simple and minimal.
|
||||||
155
local_rag.py
Normal file
155
local_rag.py
Normal file
@ -0,0 +1,155 @@
|
|||||||
|
"""
|
||||||
|
Local RAG setup with LangChain, Ollama, and FAISS
|
||||||
|
Minimal dependencies, simple code
|
||||||
|
"""
|
||||||
|
import os
|
||||||
|
from pathlib import Path
|
||||||
|
from langchain_community.document_loaders import PyPDFLoader, TextLoader
|
||||||
|
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||||||
|
from langchain_community.embeddings import HuggingFaceEmbeddings
|
||||||
|
from langchain_community.vectorstores import FAISS
|
||||||
|
from langchain_ollama import ChatOllama
|
||||||
|
|
||||||
|
|
||||||
|
class LocalRAG:
|
||||||
|
def __init__(self, vectorstore_path="./vectorstore", ollama_model="llama2"):
|
||||||
|
"""Initialize local RAG system"""
|
||||||
|
self.vectorstore_path = vectorstore_path
|
||||||
|
self.ollama_model = ollama_model
|
||||||
|
|
||||||
|
# Embeddings
|
||||||
|
print("Loading embeddings model...")
|
||||||
|
self.embeddings = HuggingFaceEmbeddings(
|
||||||
|
model_name="sentence-transformers/all-MiniLM-L6-v2"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Text splitter
|
||||||
|
self.text_splitter = RecursiveCharacterTextSplitter(
|
||||||
|
chunk_size=1000,
|
||||||
|
chunk_overlap=200
|
||||||
|
)
|
||||||
|
|
||||||
|
# Ollama LLM
|
||||||
|
print(f"Connecting to Ollama (model: {ollama_model})...")
|
||||||
|
self.llm = ChatOllama(
|
||||||
|
model=ollama_model,
|
||||||
|
base_url="http://localhost:11434"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Vector store (load if exists, otherwise None)
|
||||||
|
self.vectorstore = None
|
||||||
|
self._load_vectorstore()
|
||||||
|
|
||||||
|
def _load_vectorstore(self):
|
||||||
|
"""Load existing vector store if available"""
|
||||||
|
index_file = os.path.join(self.vectorstore_path, "index.faiss")
|
||||||
|
if os.path.exists(index_file):
|
||||||
|
try:
|
||||||
|
self.vectorstore = FAISS.load_local(
|
||||||
|
self.vectorstore_path,
|
||||||
|
self.embeddings,
|
||||||
|
allow_dangerous_deserialization=True
|
||||||
|
)
|
||||||
|
print(f"Loaded existing vector store from {self.vectorstore_path}")
|
||||||
|
except Exception as e:
|
||||||
|
print(f"Could not load vector store: {e}")
|
||||||
|
self.vectorstore = None
|
||||||
|
|
||||||
|
def add_documents(self, file_paths):
|
||||||
|
"""Add documents to the vector store"""
|
||||||
|
print(f"\nLoading {len(file_paths)} document(s)...")
|
||||||
|
all_docs = []
|
||||||
|
|
||||||
|
for file_path in file_paths:
|
||||||
|
path = Path(file_path)
|
||||||
|
if not path.exists():
|
||||||
|
print(f"Warning: {file_path} not found, skipping")
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Load document
|
||||||
|
if path.suffix.lower() == '.pdf':
|
||||||
|
loader = PyPDFLoader(str(path))
|
||||||
|
elif path.suffix.lower() in ['.txt', '.md']:
|
||||||
|
loader = TextLoader(str(path))
|
||||||
|
else:
|
||||||
|
print(f"Warning: Unsupported file type {path.suffix}, skipping")
|
||||||
|
continue
|
||||||
|
|
||||||
|
docs = loader.load()
|
||||||
|
chunks = self.text_splitter.split_documents(docs)
|
||||||
|
all_docs.extend(chunks)
|
||||||
|
print(f" - {path.name}: {len(chunks)} chunks")
|
||||||
|
|
||||||
|
if not all_docs:
|
||||||
|
print("No documents loaded!")
|
||||||
|
return
|
||||||
|
|
||||||
|
# Create or update vector store
|
||||||
|
print(f"\nCreating embeddings for {len(all_docs)} chunks...")
|
||||||
|
if self.vectorstore is None:
|
||||||
|
self.vectorstore = FAISS.from_documents(all_docs, self.embeddings)
|
||||||
|
else:
|
||||||
|
new_store = FAISS.from_documents(all_docs, self.embeddings)
|
||||||
|
self.vectorstore.merge_from(new_store)
|
||||||
|
|
||||||
|
# Save
|
||||||
|
os.makedirs(self.vectorstore_path, exist_ok=True)
|
||||||
|
self.vectorstore.save_local(self.vectorstore_path)
|
||||||
|
print(f"Vector store saved to {self.vectorstore_path}")
|
||||||
|
|
||||||
|
def query(self, question, k=4):
|
||||||
|
"""Query the RAG system"""
|
||||||
|
if self.vectorstore is None:
|
||||||
|
return "Error: No documents loaded. Please add documents first."
|
||||||
|
|
||||||
|
print(f"\nSearching for relevant documents...")
|
||||||
|
docs = self.vectorstore.similarity_search(question, k=k)
|
||||||
|
print(f"Found {len(docs)} relevant documents")
|
||||||
|
|
||||||
|
# Combine context from documents
|
||||||
|
context = "\n\n".join([doc.page_content for doc in docs])
|
||||||
|
|
||||||
|
# Create prompt
|
||||||
|
prompt = f"""Use the following context to answer the question.
|
||||||
|
If you don't know the answer, say that you don't know instead of making up an answer.
|
||||||
|
|
||||||
|
Context:
|
||||||
|
{context}
|
||||||
|
|
||||||
|
Question: {question}
|
||||||
|
|
||||||
|
Answer:"""
|
||||||
|
|
||||||
|
print("Generating answer with Ollama...")
|
||||||
|
response = self.llm.invoke(prompt)
|
||||||
|
answer = response.content if hasattr(response, 'content') else str(response)
|
||||||
|
|
||||||
|
return answer
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
"""Example usage"""
|
||||||
|
print("=" * 60)
|
||||||
|
print("Local RAG with LangChain, Ollama, and FAISS")
|
||||||
|
print("=" * 60)
|
||||||
|
|
||||||
|
# Initialize
|
||||||
|
rag = LocalRAG(ollama_model="llama2")
|
||||||
|
|
||||||
|
# Add documents (uncomment and add your file paths)
|
||||||
|
rag.add_documents([
|
||||||
|
"diverses/local_rag/test1.pdf",
|
||||||
|
"diverses/local_rag/test2.txt"
|
||||||
|
])
|
||||||
|
|
||||||
|
# Query
|
||||||
|
# question = "What is this document about?"
|
||||||
|
# answer = rag.query(question)
|
||||||
|
# print(f"\nQuestion: {question}")
|
||||||
|
# print(f"Answer: {answer}")
|
||||||
|
|
||||||
|
print("\nSetup complete! Uncomment the code above to add documents and query.")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
5
requirements.txt
Normal file
5
requirements.txt
Normal file
@ -0,0 +1,5 @@
|
|||||||
|
langchain
|
||||||
|
langchain-ollama
|
||||||
|
faiss-cpu
|
||||||
|
sentence-transformers
|
||||||
|
pypdf
|
||||||
Loading…
x
Reference in New Issue
Block a user