Retrieval-augmented generation

pip install openai requests rank_bm25
export MUNITO_API_KEY="munitoai_..."

# pip install openai requests rank_bm25
import os
import re
import sys
import time

import requests
from openai import OpenAI
from rank_bm25 import BM25Okapi

HOST = "https://api.munito.ai"
BASE_URL = f"{HOST}/inference/v1"
HEADERS = {"Authorization": f"Bearer {os.environ['MUNITO_API_KEY']}"}
DOCUMENTS = {
    "transformer": "https://arxiv.org/pdf/1706.03762",
}


def ocr_pages(url: str) -> list[str]:
    """Read a PDF by URL with Munito OCR. Return the Markdown of each page."""
    response = requests.post(
        f"{BASE_URL}/ocr",
        headers=HEADERS,
        json={"model": "munito/ocr-auto", "document": {"type": "url", "url": url}},
        timeout=300,
    )
    response.raise_for_status()
    body = response.json()
    while body.get("status") in ("accepted", "running"):
        time.sleep(2)
        poll = requests.get(HOST + body["location"], headers=HEADERS, timeout=60)
        poll.raise_for_status()
        body = {**poll.json(), "location": body["location"]}
    return [page["markdown"] for page in body["pages"]]


def chunks(doc: str, pages: list[str], size: int = 800, overlap: int = 200) -> list[dict]:
    """Cut each page into overlapping chunks of characters. Keep the source of each chunk."""
    out = []
    for number, text in enumerate(pages, start=1):
        for start in range(0, max(len(text), 1), size - overlap):
            piece = text[start : start + size].strip()
            if piece:
                out.append({"id": f"{doc} p.{number}", "text": piece})
    return out


def tokens(text: str) -> list[str]:
    return re.findall(r"\w+", text.lower())


def answer(question: str, index: BM25Okapi, corpus: list[dict], k: int = 5) -> str:
    """Retrieve the k best chunks and answer only from them, with citations."""
    scores = index.get_scores(tokens(question))
    best = sorted(range(len(corpus)), key=lambda i: scores[i], reverse=True)[:k]
    context = "\n\n".join(f"[{corpus[i]['id']}]\n{corpus[i]['text']}" for i in best)
    client = OpenAI(base_url=BASE_URL, api_key=os.environ["MUNITO_API_KEY"])
    completion = client.chat.completions.create(
        model="qwen/qwen3.8-27b",
        messages=[
            {
                "role": "system",
                "content": "Answer from the sources only. After each claim, cite its source in brackets, "
                "such as [transformer p.3]. If the sources do not answer, say so.",
            },
            {"role": "user", "content": f"Sources:\n{context}\n\nQuestion: {question}"},
        ],
        reasoning_effort="low",
        max_tokens=2048,
    )
    return completion.choices[0].message.content


if __name__ == "__main__":
    corpus = [chunk for doc, url in DOCUMENTS.items() for chunk in chunks(doc, ocr_pages(url))]
    index = BM25Okapi([tokens(c["text"]) for c in corpus])
    question = sys.argv[1] if len(sys.argv) > 1 else "Which optimizer trains the model, and how many warmup steps does it use?"
    print(answer(question, index, corpus))

python rag.py "Which optimizer trains the model, and how many warmup steps does it use?"

The model is trained using the Adam optimizer (with β1 = 0.9, β2 = 0.98, and ϵ = 10⁻⁹) [transformer p.7].
It uses 4000 warmup steps [transformer p.7].