Skip to content

Document RAG

Open In Colab

Document RAG (Retrieval-Augmented Generation)

This example demonstrates how to build a classic RAG system using Synalinks. RAG combines document retrieval with language model generation to answer questions based on your own documents.

How RAG Works

graph LR
    subgraph Indexing
        A[Documents] --> B[Embeddings]
        B --> C[(KnowledgeBase)]
    end
    subgraph Query Time
        D[Question] --> E[RetrieveKnowledge]
        C --> E
        E --> F[Relevant Docs]
        F --> G[Generator]
        G --> H[Answer]
    end
  1. Index: Store documents in a knowledge base with embeddings
  2. Retrieve: When a question is asked, find relevant documents
  3. Generate: Use the retrieved context to generate an accurate answer

Creating a Document Store

class Document(synalinks.DataModel):
    id: str = synalinks.Field(description="Document ID")
    title: str = synalinks.Field(description="Document title")
    content: str = synalinks.Field(description="Document content")

knowledge_base = synalinks.KnowledgeBase(
    uri="duckdb://./documents.db",
    data_models=[Document],
    embedding_model=embedding_model,  # For semantic search
)

Building the RAG Pipeline

inputs = synalinks.Input(data_model=Query)

# Retrieve relevant documents
retrieved = await synalinks.RetrieveKnowledge(
    knowledge_base=knowledge_base,
    language_model=language_model,
    search_type="hybrid",
    k=3,
)(inputs)

# Generate answer from retrieved context
answer = await synalinks.Generator(
    data_model=Answer,
    language_model=language_model,
    instructions="Answer based on the retrieved documents.",
)(retrieved)

Key Takeaways

  • Hybrid Search: Combines keyword (BM25) and semantic (vector) search for better retrieval accuracy.
  • Chunking: For large documents, split into smaller chunks for better retrieval granularity.
  • Context Window: Retrieved documents are passed as context to the LM for grounded generation.
  • Trainable: The retrieval and generation modules can be optimized using Synalinks training.

Program Visualization

document_rag

API References

Answer

Bases: DataModel

An answer generated from retrieved documents.

Source code in examples/13_document_rag.py
class Answer(synalinks.DataModel):
    """An answer generated from retrieved documents."""

    answer: str = synalinks.Field(
        description="The answer to the question based on retrieved documents",
    )
    sources: str = synalinks.Field(
        description="The document titles used to generate the answer",
    )

Document

Bases: DataModel

A document stored in the knowledge base.

Source code in examples/13_document_rag.py
class Document(synalinks.DataModel):
    """A document stored in the knowledge base."""

    id: str = synalinks.Field(
        description="Unique document identifier",
    )
    title: str = synalinks.Field(
        description="Document title",
    )
    content: str = synalinks.Field(
        description="The main text content of the document",
    )
    source: str = synalinks.Field(
        description="Source or category of the document",
    )

Query

Bases: DataModel

A user question.

Source code in examples/13_document_rag.py
class Query(synalinks.DataModel):
    """A user question."""

    query: str = synalinks.Field(
        description="The user's question",
    )

Source

--8 < --"examples/13_document_rag.py:89"

Run log

The log below is the unedited combined output of running the example above with local models (ollama).

Full run log: examples/13_document_rag.log
Step 1: Creating Document Knowledge Base
============================================================
Knowledge base created: ./examples/documents.db

Step 2: Adding Documents to Knowledge Base
============================================================

Storing documents...
  - Stored: Introduction to Machine Learning
  - Stored: Types of Machine Learning
  - Stored: Neural Networks Explained
  - Stored: Natural Language Processing Overview
  - Stored: Large Language Models
  - Stored: RAG: Retrieval-Augmented Generation
  - Stored: Vector Databases and Embeddings

Total documents stored: 7

Step 3: Building RAG Pipeline
============================================================
RAG pipeline created successfully!

Step 4: Testing RAG System
============================================================

────────────────────────────────────────────────────────────
Q: What is machine learning?
────────────────────────────────────────────────────────────

A: Machine learning is a subset of artificial intelligence that enables systems to learn and improve from experience without being explicitly programmed. It focuses on developing algorithms that can access data and use it to learn for themselves. The process begins with observations or data, such as examples, direct experience, or instruction, to look for patterns in data and make better decisions in the future.

Sources: AI Fundamentals

────────────────────────────────────────────────────────────
Q: What are the three types of machine learning?
────────────────────────────────────────────────────────────

A: The three main types of machine learning are supervised learning, unsupervised learning, and reinforcement learning.

Sources: ][

────────────────────────────────────────────────────────────
Q: How do neural networks work?
────────────────────────────────────────────────────────────

A: Neural networks are computing systems inspired by the human brain and consist of interconnected nodes (neurons) organized in layers: an input layer, one or more hidden layers, and an output layer. Each connection has a weight that adjusts as learning proceeds. Deep learning uses neural networks with many hidden layers to learn complex patterns. Popular architectures include CNNs for images and RNNs for sequential data.

Sources: ['Neural Networks Explained', 'Deep Learning']

────────────────────────────────────────────────────────────
Q: What is RAG and why is it useful?
────────────────────────────────────────────────────────────

A: RAG stands for Retrieval-Augmented Generation. It is a technique that combines information retrieval with text generation, reducing hallucinations and allowing for up-to-date information and domain-specific applications. This approach is useful as it enables systems to learn and improve from both the knowledge embedded in model weights and relevant documents from a knowledge base.

Sources: ['Advanced AI']

────────────────────────────────────────────────────────────
Q: What are vector databases used for?
────────────────────────────────────────────────────────────

A: Vector databases are used for storing and indexing high-dimensional vectors (embeddings) for efficient similarity search. Embeddings are numerical representations of text, images, or other data that capture semantic meaning. Similar items have vectors that are close together in the embedding space.

Sources: ['Document with title: Vector Databases and Embeddings', 'Document with title: RAG: Retrieval-Augmented Generation']

────────────────────────────────────────────────────────────
Q: What is the capital of France?
────────────────────────────────────────────────────────────

A: I don't have information about that.

Sources: ][


Step 5: Direct Search Examples
============================================================

Full-text search for 'transformer':
  - Document ID: doc-004, Score: 1.0000
  - Document ID: doc-005, Score: 1.0000

Hybrid search for 'how computers learn from data':
  - Document ID: doc-001, Score: 0.0328
  - Document ID: doc-002, Score: 0.0320
  - Document ID: doc-003, Score: 0.0320


Step 6: All Documents in Knowledge Base
============================================================

[doc-001] Introduction to Machine Learning
   Source: AI Fundamentals
   Content: Machine learning is a subset of artificial intelligence that enables             systems to learn an...

[doc-002] Types of Machine Learning
   Source: AI Fundamentals
   Content: There are three main types of machine learning: supervised learning,             unsupervised learni...

[doc-003] Neural Networks Explained
   Source: Deep Learning
   Content: Neural networks are computing systems inspired by biological neural             networks in the huma...

[doc-004] Natural Language Processing Overview
   Source: NLP Guide
   Content: Natural Language Processing (NLP) is a field of AI that focuses on             the interaction betwe...

[doc-005] Large Language Models
   Source: NLP Guide
   Content: Large Language Models (LLMs) are neural networks trained on massive             amounts of text data...

[doc-006] RAG: Retrieval-Augmented Generation
   Source: Advanced AI
   Content: Retrieval-Augmented Generation (RAG) is a technique that combines             information retrieval ...

[doc-007] Vector Databases and Embeddings
   Source: Advanced AI
   Content: Vector databases store and index high-dimensional vectors (embeddings)             for efficient sim...


Done!