Міні-пайплайн RAG з 90 рэядкамі Python: адзысканне даных, ўдосконаленне і генераванне.

Retrieval-augmented generation не ўскладненне якога-небудзь модэлю. Гэта проблема пакавання. Акна кантэксту маленькае, документы вельмі вялікія, і трэба якось выбраць, калікі фрагменты будуць включаны. Гэта рашэнне можна напісаць за адно 90 ліній простага Python, без LangChain і без службы вектарных баз дадзеных. У гэтага працэсу тры крокі:

1. Працэўваць з дадзенымі. Увёсці запіт, ранжаваць фрагменты за сэмплівасцю і застаўці топ-k.
2. Дапамагаць сабе дадзенымі. З’еднаныя фрагменты включыць у сістэмны запіс.
3. Ствараць выходны тэкст. Выкарыстоўваць той жа прыем чату, які і ў будзь-ям базовам скрыпце LLM.

Кожная абстракцыя, яку пропонуе фрэймворк урок LangChain, пасля ператвараецца ў адну з гэтых трох катэгорыяў, таму спачатку разуменне простай версіі дапамагае лёгкаярозумеласці фрэймворку.

Увёсцы ператвараюць значэнне ў геаметрыю

Мадэль, якая стварае эмбедынгі, ператварае фрагмент тэксту на вектор – тое ўзроўню длігага раду чысла з плаваючай канмою, распакаванага так, каб тэксты з супадзячым значэнням належалі блізка адзін да другога. Адстань вычысляецца па косінусовай схожасці: 1,0 для ідэнтычнага направлення, 0,0 для вялікай розніцы. У гэтым і заключаецца весь секрэт. Калі значэння стае геаметрыяй, пошук рэлевантных дакументаў ператвараецца на пошук вектораў, якія знаходзяцца блізка.

# pip install openai
# export OPENAI_API_KEY=sk-...
import math

from openai import OpenAI

client = OpenAI()


def embed(texts):
    resp = client.embeddings.create(
        model="text-embedding-3-small",
        input=texts,
    )
    return [d.embedding for d in resp.data]


def cosine(a, b):
    dot = sum(x * y for x, y in zip(a, b))
    na = math.sqrt(sum(x * x for x in a))
    nb = math.sqrt(sum(y * y for y in b))
    return dot / (na * nb)


q, hit, miss = embed([
    "How do I get my money back?",
    "Refunds are possible within 14 days "
    "of purchase.",
    "Data is stored in Frankfurt and "
    "encrypted at rest.",
])

print(f"refund doc:  {cosine(q, hit):.3f}")
print(f"storage doc: {cosine(q, miss):.3f}")

Тыповым рэзультатам ад гэтага скрыпта ёсць значэнне 0.531 для дакумента пра вярненне грошаў і 0.126 для дакумента пра зберагачча. Паглядзіце, як адбываецца паўтарэнне: у запытце напісана «вярненне грошаў», а ў дакуменце — «працэсы вярнення». Не існуе спаканых ключоўых словаў, пры тым сэрабістас вялікі. Самэ гэтаму вы вікорыстоўваеце вбудоўванне заместо пошуку.

Вектарны зберагач — гэта спіс, які можна рангаваць

«Вектарны зберагач» звучыць як інфраструктура. У малых масштабах гэта спіс пар (вектар, тэкст) і спосаб іх сортавання. Pinecone, Chroma і pgvector працуюць з мільйонамі вектараў, але гэтычна важны інтэрфейс, і ён падпадае пад адну калякцію:

class VectorStore:
    def __init__(self):
        self.texts = []
        self.vectors = []

    def add(self, texts):
        self.texts.extend(texts)
        self.vectors.extend(embed(texts))

    def top_k(self, query, k=3):
        q = embed([query])[0]
        scored = [
            (cosine(q, v), t)
            for v, t in zip(
                self.vectors, self.texts
            )
        ]
        scored.sort(
            key=lambda s: s[0], reverse=True
        )
        return scored[:k]


store = VectorStore()
store.add([
    "Refunds are possible within 14 days.",
    "The Pro plan costs 12 EUR per month.",
    "Data is stored in Frankfurt.",
])

for score, text in store.top_k(
    "can I cancel and get a refund?", k=2
):
    print(f"{score:.3f}  {text}")

top_k — гэта крытэрыя для выбору таго, што будзе застаўлены, і ён вядомы: усё ранжуецца, застаўляюцца k элементаў, а рэшта адхоўляецца. Значэнне k — гэта «рэгулятор бюджэту», які дазволяе сабраныць больш інформацыі за рахунак месца ў вікні кантэксту і колькасці токенав, якія платаюцца.

Адзінаванне функцый збору, дапамогі і стварэння

Тут новая толькі адна функцыя — build_messages. Яна прызначае номеры паводлі рангу, дадвае іх у сістэмныя запыты і прасіць мадэль адпаведзаць толькі на ўсё, што было дана. У запыце корыстніка ўказваецца толькі запітанне. Спіс запытоў такі ж, як у будзь-ям чаце, таму RAG зменяе толькі ўсё, што ў ям ёсць. Шырокавядомы файл працюе без змян:

# examples/part07_rag.py
# Mini RAG in ~90 lines: embed -> store -> top-k
# -> stitch into messages -> generate.
#
# Setup:
#   pip install openai
#   export OPENAI_API_KEY=sk-...
#
# Run:
#   python part07_rag.py

import math

from openai import OpenAI

client = OpenAI()

EMBED_MODEL = "text-embedding-3-small"
CHAT_MODEL = "gpt-4o-mini"

# A tiny corpus. In real life these are chunks
# of your docs, not whole documents.
DOCS = [
    "Support hours are 9am to 6pm CET, "
    "Monday to Friday.",
    "Refunds are possible within 14 days "
    "of purchase, no questions asked.",
    "The Pro plan costs 12 EUR per month "
    "and includes API access.",
    "The Free plan allows 3 projects and "
    "community support only.",
    "Data is stored in Frankfurt and "
    "encrypted at rest with AES-256.",
    "You can export all your data as "
    "JSON from the settings page.",
]


def embed(texts):
    resp = client.embeddings.create(
        model=EMBED_MODEL,
        input=texts,
    )
    return [d.embedding for d in resp.data]


def cosine(a, b):
    dot = sum(x * y for x, y in zip(a, b))
    na = math.sqrt(sum(x * x for x in a))
    nb = math.sqrt(sum(y * y for y in b))
    return dot / (na * nb)


class VectorStore:
    """A list of (vector, text) pairs. That is
    all a vector store is, before you need
    scale."""

    def __init__(self):
        self.texts = []
        self.vectors = []

    def add(self, texts):
        self.texts.extend(texts)
        self.vectors.extend(embed(texts))

    def top_k(self, query, k=3):
        q = embed([query])[0]
        scored = [
            (cosine(q, v), t)
            for v, t in zip(
                self.vectors, self.texts
            )
        ]
        scored.sort(key=lambda s: s[0], reverse=True)
        return scored[:k]



def build_messages(question, hits):
    context = "\n\n".join(
        f"[{i + 1}] {text}"
        for i, (_, text) in enumerate(hits)
    )
    system = (
        "Answer using only the context below. "
        "If the answer is not in the context, "
        "say you do not know.\n\n"
        "Context:\n" + context
    )
    return [
        {"role": "system", "content": system},
        {"role": "user", "content": question},
    ]


def ask(store, question):
    hits = store.top_k(question, k=3)
    print(f"Q: {question}")
    for score, text in hits:
        print(f"  {score:.3f}  {text[:44]}...")
    messages = build_messages(question, hits)
    resp = client.chat.completions.create(
        model=CHAT_MODEL,
        messages=messages,
    )
    return resp.choices[0].message.content


if __name__ == "__main__":
    store = VectorStore()
    store.add(DOCS)

    answer = ask(
        store,
        "Where is my data stored, and can "
        "I get it out?",
    )
    print(f"A: {answer}")

У прыёмнай працэ спытваец уточняе, дзе зберагаюцца даны корыстніка і чы можна іх экспортуваць. Тры найлепшыя рэзультаты: 0,588 за частку пра зберагачча даных у Франкфурце з шыфраваннем, 0,472 за частку пра экспорт усіх даных у фармаце JSON і 0,244 за правіла вярнення грошаў. Мадэль адпавядае, што даны знаходзяцца у Франкфурце, шыфруюцца па час зберагачча за дапамою AES-256, і што на сторонцы налашоўкі ёсць можлівасць повнага экспорту у фармаце JSON.

Обе неабходныя часткі — пра зберагачча і пра экспорт — былі прыняты. Частка пра вярнення грошаў зайняла трэцяе месца, але мадэль яе проігноравала. Адзысканне інформацыі можа быць непаштоўным, як толькі правільны матэрыял опынуцца вышэй у рэйтынгу.

Дзе ламаеця прыбліжная версія

У этых 90 рэчэй ёсць чыстыя мерыябы, і кожная з іх паказвае на тое, што дадаюць справжнія сістэмы:

• Чаканаванне. У дэме зберагаюцца цэлыя рэчы. Настоямыя дакументы трэба раздзеляць з перакрыццю, і гэта ўхідна праця, а не рашэнне мадэлі.
• Павторачунак укладання. Укладанне запіту з кожнага вызову ўсё правільна. Але павторачунак укладання всего корпуса з кожнага пачатку работы — ні, таму трэба зберагаць вектары.
• Лінейны сканаванне. Праця з косай сэмабільнасцю з кожным вектарам выконваецца за час O(n). Пасля прыблізна 100 тысяч чаканаў стае цікавым індэks на алгорытмах нейронных сетак, і самэ гэта робіць вектарныя базы дадзеных практычнымі.
• Якасць ранжавання. Простая косая метода top-k ёсць прымітным інструментам, а ранжавальнікі чы ўзаўмешаныя пошукавыя сістэмы робяць яе болей точной.

LangChain абгружае кожны з гэтых элементаў як компонент, які можна зменіць. Калі первісны цыкл запішаны вручную, гэтыя компоненты выглядаюць такімі, якімі ёсць, — толькі зручнасцямі, без якіх-небудзь містычных элементаў.