Skip to content

RAG over Research Papers

Build a Retrieval-Augmented Generation (RAG) pipeline that lets you ask questions over a collection of research PDFs. Uses PAIS embeddings to index documents and the PAIS LLM to generate answers.

What you'll build

flowchart LR
    PDFs[PDF Papers] --> Chunk[Chunk & Embed\nPAIS embeddings]
    Chunk --> VDB[(Chroma\nvector store)]
    Q[Research\nquestion] --> Retrieve[Retrieve top-k\nrelevant chunks]
    VDB --> Retrieve
    Retrieve --> LLM[PAIS LLM\nLlama 3.1 8B]
    LLM --> A[Grounded answer\nwith citations]

Prerequisites

pip install llama-index llama-index-llms-openai-like llama-index-embeddings-openai \
            llama-index-vector-stores-chroma chromadb pypdf python-dotenv

Full example

# rag_pipeline.py
"""
RAG pipeline over a PDF paper collection using PAIS + LlamaIndex.

Usage:
    python rag_pipeline.py --papers ./papers --query "What methods are used?"
"""
import os
import argparse
from pathlib import Path
from dotenv import load_dotenv

import chromadb
from llama_index.core import (
    Settings,
    SimpleDirectoryReader,
    VectorStoreIndex,
    StorageContext,
)
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.postprocessor import SimilarityPostprocessor
from llama_index.llms.openai_like import OpenAILike
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.vector_stores.chroma import ChromaVectorStore

load_dotenv()

# ── Configure PAIS models ──────────────────────────────────────────────────────

def configure_pais():
    Settings.llm = OpenAILike(
        model="llama-3-1-8b",
        api_base=os.environ["PAIS_API_BASE"],
        api_key=os.environ["PAIS_API_KEY"],
        is_chat_model=True,
        context_window=128_000,
        max_tokens=1024,
        temperature=0.1,
    )
    Settings.embed_model = OpenAIEmbedding(
        model="qwen3-vl-embedding-8b",
        api_base=os.environ["PAIS_API_BASE"],
        api_key=os.environ["PAIS_API_KEY"],
    )
    Settings.chunk_size = 512
    Settings.chunk_overlap = 64


# ── Build or load index ────────────────────────────────────────────────────────

def build_index(papers_dir: str, persist_dir: str = "./chroma_db") -> VectorStoreIndex:
    chroma_client = chromadb.PersistentClient(path=persist_dir)
    collection = chroma_client.get_or_create_collection(
        "research_papers",
        metadata={"hnsw:space": "cosine"},
    )
    vector_store = ChromaVectorStore(chroma_collection=collection)
    storage_context = StorageContext.from_defaults(vector_store=vector_store)

    # If collection already has documents, load the existing index
    if collection.count() > 0:
        print(f"Loading existing index ({collection.count()} chunks)...")
        return VectorStoreIndex.from_vector_store(
            vector_store, storage_context=storage_context
        )

    print(f"Building index from {papers_dir}...")
    documents = SimpleDirectoryReader(
        papers_dir,
        recursive=True,
        required_exts=[".pdf", ".txt", ".md"],
    ).load_data()
    print(f"Loaded {len(documents)} document sections")

    index = VectorStoreIndex.from_documents(
        documents,
        storage_context=storage_context,
        show_progress=True,
    )
    print(f"Index built: {collection.count()} chunks stored")
    return index


# ── Query with citations ───────────────────────────────────────────────────────

def query_with_citations(index: VectorStoreIndex, question: str, top_k: int = 5) -> None:
    retriever = VectorIndexRetriever(index=index, similarity_top_k=top_k)
    postprocessor = SimilarityPostprocessor(similarity_cutoff=0.3)

    query_engine = RetrieverQueryEngine.from_args(
        retriever=retriever,
        node_postprocessors=[postprocessor],
        response_mode="compact",
    )

    print(f"\nQuestion: {question}\n")
    response = query_engine.query(question)

    print("Answer:")
    print(response.response)

    print("\nSources:")
    seen = set()
    for node in response.source_nodes:
        fname = node.metadata.get("file_name", "unknown")
        if fname not in seen:
            seen.add(fname)
            page = node.metadata.get("page_label", "?")
            print(f"  [{node.score:.3f}] {fname} (p. {page})")


# ── Main ───────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--papers", default="./papers", help="Directory of PDFs")
    parser.add_argument("--query", required=True, help="Question to answer")
    parser.add_argument("--top-k", type=int, default=5)
    parser.add_argument("--rebuild", action="store_true", help="Rebuild index from scratch")
    args = parser.parse_args()

    configure_pais()

    persist_dir = "./chroma_db"
    if args.rebuild and Path(persist_dir).exists():
        import shutil
        shutil.rmtree(persist_dir)
        print("Removed existing index")

    index = build_index(args.papers, persist_dir)
    query_with_citations(index, args.query, args.top_k)


if __name__ == "__main__":
    main()

Running the pipeline

# First run — builds the index (calls embedding API for all PDFs)
python rag_pipeline.py --papers ./papers \
    --query "What hyperparameter optimisation methods are compared?"

# Subsequent runs — loads existing index, fast
python rag_pipeline.py --query "What datasets are used across these papers?"

# Rebuild index after adding new papers
python rag_pipeline.py --papers ./papers --rebuild \
    --query "Which papers use transformer architectures?"

Example output

Building index from ./papers...
Loaded 847 document sections
Index built: 1203 chunks stored

Question: What hyperparameter optimisation methods are compared?

Answer:
The papers compare Bayesian optimisation (primarily TPE via Optuna), random search,
and grid search for hyperparameter tuning. Li et al. (2023) find that Bayesian
optimisation converges 3× faster than random search on large language model
fine-tuning tasks. Several papers also evaluate population-based training (PBT)
for continuous adaptation of learning rate schedules.

Sources:
  [0.847] li_2023_hpo_survey.pdf (p. 4)
  [0.821] zhang_hyperopt_transformers.pdf (p. 2)
  [0.794] automl_benchmark_2024.pdf (p. 8)

Extending the pipeline

Add a reranker for improved retrieval precision: 

from llama_index.postprocessor.cohere_rerank import CohereRerank
# Or use a cross-encoder reranker served on PAIS (future capability)

Multi-index query over different collections (e.g., papers + your own notes): 

from llama_index.core.indices.composability import ComposableGraph

Streaming responses for interactive use: 

streaming_engine = index.as_query_engine(streaming=True)
response = streaming_engine.query("Summarise the key findings...")
response.print_response_stream()

See also