Vlastní data (RAG) pro vašeho AI agenta část první - krok za krokem s Python a PostgreSQL

Srovnejte to s tím, když vám řeknu, že si klidně můžete před odpovědí pohledat po Internetu, prolistovat poznámky, prolétnout vytištěné směrnice nebo si zaklikat v Azure dokumentaci. Předpokládám, že tím se jednak výrazně zlepší přesnost vašich odpovědí a také mi řeknete, kde se o tom píše, ocitujete zdroje.

Jenže co když místo dobrého zacílení na určitou část dokumentace nebo směrnici vám nabídnu stránky z celé knihovny, ale bez přebalů a jakéhokoli řazení nebo indexování titulů? Pomůže to? Asi ano, ale bude vám to trvat mnohem déle, spotřebujete daleko víc energie a může se stát, že pokud naprostá většina stránek není pro odpověď důležitá, tak ztratíte pozornost a správné odpovědi si tam někde uprostřed nevšimnete.

To, co jsem právě popsal, funguje ve velkých jazykových modelech (LLM) taky a říká se tomu Retrieval Augmented Generation (RAG).

Ideální je, když se model dostane k co nejlépe zacíleným dokumentům, ve kterých se skrývá odpověď, a on tak může dávat naprosto přesné, rychlé a levné výstupy. Když se netrefíme a dáme mu dokumenty, které moc relevantní nejsou, odpovědi se zhorší. Když budeme mít model s velkým kontextem a nahrneme mu celou knihovnu, bude to drahé, pomalé a výsledky stejně nebudou tak dobré, jako v případě dobře zvolených referenčních dokumentů.

Pracovat budeme s PostgreSQL, protože se tam můžu naučit, co se všechno děje. Pro svoje projekty ale určitě zvažte i Azure AI Search, který tohle všechno rovnou jednoduše umí a kromě toho řeší i problematiku krájení a nabírání.

Všechny techniky si ukážeme na jednoduchém příkladu s filmy a jejich popisky. Příště se pak pustíme do složitějších technik vyžadujících hlubší zpracování dat, hierarchické a grafové přístupy a sumarizace.

Otevřete si prosím Jupyter Notebook na mém GitHubu. Najdete tam konkrétní výsledky a výstupy, doporučuji si ho vzít při čtení k ruce.

Celý příběh začíná tím, že vezmeme dokumenty, kterými budou popisky filmů, a vložíme je do PostgreSQL s tím, že vytvoříme kombinovaný sloupeček s názvem a popisem (ať se mi dobře dělají vektory). Potřebnou infrastrukturu jako jsou Azure OpenAI modely a Azure PostgreSQL Flexible Server a vše si můžete nahodit z mého Terraformu v příslušném adresáři.

Než udělám samotný import, je potřeba přidat některá rozšíření:

Rozšíření	K čemu ho používám
azure_ai	Přímo z PSQL provolávám Azure OpenAI pro embeddingy a případně další Azure AI služby.
pgvector	Ukládání embeddingů přímo do PostgreSQL a vektorové vyhledávání.
pg_diskann	Rychlejší a efektivnější vyhledávání v embeddingových vektorech.

Založím tabulku, přidám sloupeček embedding s 2000 dimenzemi, přidám další sloupeček full_text_search pro tsvector (statistika pro full text, víceméně kořeny použitých slov a jejich výskyty) a přidám DiskANN index. Pak vložím data a následně k nim přidám embeddings s využitím SQL příkazu UPDATE společně s funkcí z azure_ai, takže databáze si do Azure OpenAI služby bude volat sama a zařídí si to.

Nejprve nasaďme full-text. PostgreSQL prošel jednotlivé popisky, identifikoval kořeny slov, spočítal si je a statistiky si u každého filmu uložil. Vzneseme dotaz, proženeme full-textem a necháme si prvních 10 nejbližších odpovědí. Jednotlivé otázky a odpovědi najdete v odkazovaném Jupyter notebooku a následně je předhodíme LLM pro formulaci odpovědi.

První pokus nic moc. Otázka na Abby vede jen na jeden film (což je málo) a ostatní také nejsou moc dobré. Potíž je, že full-text search není ChatGPT, takže otázka What movies are about Abby? je značně nevhodná, protože tím hledáme všechna tato slova, z nichž kromě Abby vlastně žádné není důležité a nepotřebujeme ho v textu vůbec mít.

Ve strojovém učení je časté, že chceme svět s jeho velkou variabilitou a složitostí zjednodušit na menší počet vlastností, defacto ho zkomprimovat do menšího a to ztrátovým způsobem. Nejde tedy o zip, ale spíš si představte něco jako MP3.

Představme si třeba, že pro text připravíme score v několika dimenzích, například „míra veršovanosti“, „míra humoru“ a tak podobně. Pokud takových vlastností použijeme stovky, budou mít dobrou vypovídací schopnost. Ve strojovém učení ale nebudeme AI diktovat, co to má být za dimenze a vlastnosti, to necháme na něm.

Tyto vektory (vyjmenovaná čísla, score za jednotlivé vlastnosti, dimenze), embeddingy, se pak používají pro reprezentaci původního vstupu a významovou extrakci u textů, obrázků i třeba difúzních modelů.

Jak v tom hledat? Nejprve ke všem popiskům filmů nechám model vytvořit embedding, ten si uložíme a až pak dostaneme dotaz, tak i ten zasadíme do latentního prostoru a koukneme, které body jsou dotazu nejpodobnější.

Třetí varianta v notebooku je sémantický search, tedy vektorové hledání podobností z embeddingů. Na první otázce, tedy filmy kde je postava Abby, nám to selhává. Tímto příkladem jsem chtěl demonstrovat, že vektory jsou povětšinou skvělé, ale v okamžiku, kdy vlastně potřebujeme hledat konkrétní krátkou faktickou informaci typu číslo faktury, tak to moc nedávají.

Na druhé dvě otázky jsou teď určitě lepší, byť Star Wars jsou stále problém. Přestože tentokrát už máme ve výsledcích i podobné sci-fi filmy, přímo těch Star Wars, co nepotřebujeme, je tam pořád dost. Je to něco, s čím v jednoduchých scénářích nic moc neuděláme - na embedding je to moc komplikovaná konstrukce.

Je zřejmé, že na některé otázky je vhodné sémantické hledání, ale jinde je úspěšnější klasický full-text. Pojďme tedy používat oboje a nasadit hybridní hledání.

SQL dotaz trochu naboptnal, ale nic hrozivého. Nejprve provedu full-text search a vrátím si 15 výsledků, následně uděláme sémantický search a dostanu také 15 výsledků. Do obou tabulek jsem si přidal pořadové číslo (index), prioritu, aby se mi to pak dobře slévalo, a taky sloupeček s použitou metodou. Následně udělám union obou tabulek a proženu to dalším dotazem, kde provedu deduplikaci podle ID filmu. V závěru takto odfiltrované výsledky seřadím podle priority a oříznu po 10 řádcích.

Hybridní hledání v PostgreSQL

WITH fulltext AS (
    SELECT 
        id, 
        title, 
        overview, 
        combined_text, 
        ts_rank(full_text_search, plainto_tsquery('english', %s)) AS score,
        'fulltext' AS method,
        ROW_NUMBER() OVER (ORDER BY ts_rank(full_text_search, plainto_tsquery('english', %s)) DESC) AS index
    FROM movies
    ORDER BY score DESC
    LIMIT 15
),
semantic AS (
    SELECT 
        id, 
        title, 
        overview, 
        combined_text, 
        (embedding <=> azure_openai.create_embeddings('text-embedding-3-large', %s, 2000, max_attempts => 5, retry_delay_ms => 500)::vector) AS score,
        'semantic' AS method,
        ROW_NUMBER() OVER (ORDER BY (embedding <=> azure_openai.create_embeddings('text-embedding-3-large', %s, 2000, max_attempts => 5, retry_delay_ms => 500)::vector)) AS index
    FROM movies
    ORDER BY score 
    LIMIT 15
),
combined_results AS (
    SELECT * FROM fulltext
    UNION ALL
    SELECT * FROM semantic
),
deduped_results AS (
    SELECT DISTINCT ON (id) 
        id, 
        title, 
        overview, 
        combined_text, 
        score, 
        method, 
        index
    FROM combined_results
)
SELECT * FROM deduped_results
ORDER BY index
LIMIT 10;

Výsledky na Abby jsou někde mezi full-textem a sémantikou, což logicky celkem odpovídá. To znamená, že jsem odstranil extrémní situace, které jsou špatné - Abby přes sémantiku a Star Wars přes full-text.

Pojďme do toho teď zamontovat další magii - re-ranking. Jde v zásadě o typicky malý jazykový model, který je specificky trénován na schopnost zhodnotit relevanci dokumentu k dané otázce.

Je příliš velký na to, abych podle něj dělal hledání o složitosti O(n), ale na druhou stranu dostatečně rychlý na to, abych pár desítek či set předvybraných výsledků podrobil zkoumání. Protože zatím není k dispozici jako extension pro PostgreSQL v Azure Flexible Server, tak to dělám na serverové straně v Pythonu a používám jednoduchý model MultiBERT, který má asi jen 33M parametrů a v pohodě poběží i v CPU.

SQL dotaz je stejný, jen si vezmu všech 30 výsledků, MultiBERTem je seřadím dle relevance a zase vezmu top 10 a pošlu do LLM. Výsledky jsou nejlepší.

RAG metody, které jsme si společně vyzkoušeli, nejsou žádná raketová věda, ale je vidět, že použité techniky vedou k značnému zlepšení výsledků. Nicméně zůstává pár problémů.

Na to se zaměříme někdy příště - zkusíme data vzít, extrahovat z nich klíčové atributy a ty mezi sebou vztahově propojit a ještě přidat sumarizační a agregované pohledy.

Ještě nutno říct, že jsem použil PostgreSQL, protože mě fascinuje jeho univerzálnost a rozšiřitelnost a je to nejlepší způsob, jak se věci naučit a zjistit, co se tam děje.

Nicméně objevuji tak trochu kolo a pokud jste v cloudu, zkuste rovnou použít Azure AI Search, který nejen co jsme dělali, ale i další techniky má rovnou v sobě. Navíc je to nejen řešení pro ukládání a hledání, ale i pro zpracování vstupů, ať už jde o indexaci, extrakci, řezání a další součásti RAG, které rozebereme také někdy příště.