Add main.py
This commit is contained in:
@@ -0,0 +1,60 @@
|
|||||||
|
"""
|
||||||
|
Minimal Retrieval‑Augmented Generation example using LangChain, Chroma and Ollama.
|
||||||
|
|
||||||
|
The script demonstrates:
|
||||||
|
1. Splitting a sample text into chunks.
|
||||||
|
2. Creating embeddings with Ollama.
|
||||||
|
3. Storing chunks in a local Chroma vector store.
|
||||||
|
4. Querying the store and generating a response with an Ollama LLM.
|
||||||
|
|
||||||
|
Prerequisites:
|
||||||
|
- Ollama must be installed and a model (e.g. tinyllama) available.
|
||||||
|
- Python dependencies from requirements.txt must be installed.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||||
|
from langchain_chroma import Chroma
|
||||||
|
from langchain_ollama import OllamaEmbeddings, ChatOllama
|
||||||
|
|
||||||
|
# 1. Sample text
|
||||||
|
SAMPLE_TEXT = (
|
||||||
|
"LangChain is a framework for developing applications powered by language models. "
|
||||||
|
"It provides abstractions for prompt construction, chaining, and memory. "
|
||||||
|
"Chroma is a fast, lightweight vector database that can be used as a backend "
|
||||||
|
"for retrieval‑augmented generation. Ollama offers locally hosted LLMs that "
|
||||||
|
"can be used for embeddings and generation.")
|
||||||
|
|
||||||
|
# 2. Split text into chunks
|
||||||
|
text_splitter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=20)
|
||||||
|
documents = text_splitter.split_text(SAMPLE_TEXT)
|
||||||
|
|
||||||
|
# 3. Create embeddings
|
||||||
|
embeddings = OllamaEmbeddings(model="tinyllama")
|
||||||
|
|
||||||
|
# 4. Create a Chroma store in a temporary directory
|
||||||
|
persist_dir = Path("./chroma_db")
|
||||||
|
persist_dir.mkdir(exist_ok=True)
|
||||||
|
|
||||||
|
vector_store = Chroma.from_texts(
|
||||||
|
texts=documents,
|
||||||
|
embedding=embeddings,
|
||||||
|
persist_directory=str(persist_dir),
|
||||||
|
)
|
||||||
|
vector_store.persist()
|
||||||
|
|
||||||
|
# 5. Query the store
|
||||||
|
query = "What is LangChain used for?"
|
||||||
|
results = vector_store.similarity_search(query, k=1)
|
||||||
|
retrieved_text = results[0].page_content
|
||||||
|
|
||||||
|
# 6. Generate a response using Ollama LLM
|
||||||
|
llm = ChatOllama(model="tinyllama")
|
||||||
|
prompt = f"Answer the following question using the provided context:\n\nContext: {retrieved_text}\n\nQuestion: {query}\nAnswer:"
|
||||||
|
response = llm.invoke(prompt)
|
||||||
|
print("\n--- Generated Response ---\n")
|
||||||
|
print(response)
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
pass
|
||||||
Reference in New Issue
Block a user