Wyszukiwarka semantyczna w Pythonie: projekt w numpy i sentence-transformers

To ostatnia część cyklu. W pierwszej nauczyłeś się liczyć embeddingi i podobieństwo cosinusowe, w drugiej dzielić dokumenty na fragmenty i zapisywać ich źródło. Teraz składamy to w mały, ale kompletny program: wyszukiwarkę, która na pytanie „Chcę odzyskać to, co zapłaciłem” znajdzie punkt regulaminu o zwrocie pieniędzy, mimo że nie padło w nim ani jedno wspólne słowo kluczowe.

Wszystko działa w pamięci komputera, bez serwera i bez bazy danych. Taki projekt jest dobrym punktem wyjścia, zanim sięgniesz po bazę wektorową.

Co zbudujemy i jak będzie wyglądał projekt?

Program ma dwie komendy: build buduje indeks z plików Markdown, a ask odpowiada na pytanie, pokazując fragmenty wraz ze źródłem. Struktura katalogu:

semsearch/
├── docs/                # dokumenty do przeszukania (.md)
├── index/               # powstaje po komendzie build
│   ├── vectors.npy      # embeddingi fragmentów
│   ├── chunks.json      # fragmenty z metadanymi
│   └── meta.json        # nazwa modelu i wymiar wektora
├── chunking.py          # z poprzedniej części
├── index.py             # budowanie indeksu
├── search.py            # wczytanie indeksu i wyszukiwanie
├── cli.py               # wiersz poleceń
├── evaluate.py          # pomiar trafności
└── tests/test_search.py

Potrzebujesz Pythona 3.10+ i trzech bibliotek:

pip install sentence-transformers numpy tiktoken pytest

Jakie dokumenty przeszukujemy?

Do testu użyliśmy trzech krótkich, fikcyjnych dokumentów o serwisie z kursami online: regulaminu z poprzedniej części oraz dwóch nowych plików. Każdy ma tytuł # i sekcje ##, bo na tej strukturze opiera się nasz podział na fragmenty. Plik docs/pomoc-techniczna.md:

# Pomoc techniczna Kursowni

## Aplikacja mobilna

Aplikacja Kursowni działa na telefonach z systemem iOS 15 lub nowszym oraz Android 10 lub nowszym. Pobierzesz ją za darmo w App Store i Google Play.

Po zalogowaniu postępy w kursach synchronizują się automatycznie między telefonem a przeglądarką.

## Wideo się nie odtwarza

Jeśli lekcja wideo się zacina, sprawdź połączenie z internetem i spróbuj zmienić jakość obrazu na niższą. Odśwież stronę lub uruchom przeglądarkę ponownie.

Gdy problem nie znika, wyłącz rozszerzenia blokujące reklamy i napisz do nas, podając nazwę przeglądarki oraz tytuł lekcji.

## Certyfikaty ukończenia

Certyfikat w formacie PDF pobierzesz po ukończeniu wszystkich lekcji i zaliczeniu quizu końcowego na co najmniej 70 procent. Znajdziesz go w zakładce „Moje certyfikaty".

## Kontakt z obsługą

Odpowiadamy na wiadomości w dni robocze od 9:00 do 17:00, zwykle w ciągu jednego dnia roboczego. Napisz na adres pomoc@kursownia.example lub użyj formularza w aplikacji.

Plik docs/kursy.md:

# Oferta kursów Kursowni

## Kursy dla początkujących

Kurs „Python od zera" obejmuje 40 lekcji i trwa około 20 godzin. Nie wymaga żadnej wiedzy wstępnej, a na końcu budujesz własną aplikację konsolową.

Kurs „Podstawy SQL" składa się z 25 lekcji i uczy zapytań do relacyjnych baz danych na przykładzie sklepu internetowego.

## Kursy dla zaawansowanych

Kurs „FastAPI w praktyce" jest przeznaczony dla osób, które znają Pythona. Przez 30 lekcji budujesz API z uwierzytelnianiem i bazą danych.

Kurs „Testowanie w pytest" pokazuje, jak pisać testy jednostkowe i integracyjne oraz jak mierzyć pokrycie kodu.

## Dostęp do materiałów

Po wykupieniu abonamentu masz dostęp do wszystkich kursów jednocześnie. Materiały możesz oglądać w dowolnym tempie, bez limitu liczby powtórzeń.

Trzeci plik, docs/regulamin.md, to regulamin z poprzedniej części. Wszystkie trzy dokumenty dają razem 19 fragmentów. Możesz też użyć własnych dokumentów (regulaminu, dokumentacji, notatek). Wyniki będą wtedy inne niż w tabelach poniżej, ale metoda ta sama.

Jak zbudować indeks z dokumentów?

Z poprzedniej części skopiuj do pliku chunking.py funkcje count_tokens, klasę Chunk oraz split_by_sections z pomocniczą _make. Budowanie indeksu to wczytanie plików, podział na fragmenty, policzenie embeddingów i zapis na dysk:

import json
from dataclasses import asdict
from pathlib import Path

import numpy as np
from sentence_transformers import SentenceTransformer

from chunking import Chunk, split_by_sections

MODEL_NAME = "paraphrase-multilingual-MiniLM-L12-v2"


def build_index(docs_dir: Path, index_dir: Path, max_tokens: int = 60) -> int:
    chunks: list[Chunk] = []
    for plik in sorted(docs_dir.glob("*.md")):
        tekst = plik.read_text(encoding="utf-8")
        chunks += split_by_sections(tekst, plik.name, max_tokens)

    model = SentenceTransformer(MODEL_NAME)
    wektory = model.encode([c.text for c in chunks], normalize_embeddings=True)

    index_dir.mkdir(exist_ok=True)
    np.save(index_dir / "vectors.npy", wektory)
    (index_dir / "chunks.json").write_text(
        json.dumps([asdict(c) for c in chunks], ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    (index_dir / "meta.json").write_text(
        json.dumps({"model": MODEL_NAME, "dim": int(wektory.shape[1])}),
        encoding="utf-8",
    )
    return len(chunks)

Kilka decyzji, które mają znaczenie:

• Model jest wielojęzyczny, bo dokumenty są po polsku (dlaczego, pokazaliśmy w pierwszej części).
• max_tokens=60 to wynik pomiaru z poprzedniej części, a nie zgadywania.
• Wektory są znormalizowane, więc iloczyn skalarny jest podobieństwem cosinusowym.
• meta.json zapisuje nazwę modelu i wymiar. Dzięki temu przy wyszukiwaniu użyjesz tego samego modelu co przy budowie indeksu, a niezgodny indeks wykryjesz od razu.

Indeks z naszych trzech dokumentów to macierz 19 × 384 w float32, czyli ok. 29 KB.

Jak wczytać indeks i znaleźć najlepsze fragmenty?

Wyszukiwanie to mnożenie macierzy wektorów przez wektor zapytania i wzięcie k największych wyników. Dodajemy do tego walidację, która wyłapuje dwa błędy z poprzednich części: zły kształt zapytania i niezgodny wymiar:

import json
from dataclasses import dataclass
from pathlib import Path

import numpy as np


@dataclass
class Index:
    vectors: np.ndarray
    chunks: list[dict]
    model_name: str


@dataclass
class Result:
    score: float
    chunk: dict


def load_index(index_dir: Path) -> Index:
    meta = json.loads((index_dir / "meta.json").read_text(encoding="utf-8"))
    vectors = np.load(index_dir / "vectors.npy")
    chunks = json.loads((index_dir / "chunks.json").read_text(encoding="utf-8"))
    if len(chunks) != len(vectors) or vectors.shape[1] != meta["dim"]:
        raise ValueError("Indeks jest uszkodzony: liczba fragmentów lub wymiar się nie zgadza.")
    return Index(vectors, chunks, meta["model"])


def search(index: Index, query_vector: np.ndarray, k: int = 3, min_score: float = 0.0) -> list[Result]:
    if query_vector.shape != (index.vectors.shape[1],):
        raise ValueError(
            f"Zły kształt zapytania {query_vector.shape}, oczekiwano ({index.vectors.shape[1]},)."
        )
    wyniki = index.vectors @ query_vector
    kolejnosc = np.argsort(-wyniki)[:k]
    return [Result(float(wyniki[i]), index.chunks[i]) for i in kolejnosc if wyniki[i] >= min_score]

Funkcja search nie ładuje modelu: dostaje gotowy wektor zapytania. Dzięki temu da się ją testować bez pobierania modelu (zrobimy to w pytest) i używać z dowolnym źródłem embeddingów.

Jak dodać wiersz poleceń?

Cienka warstwa na wierzchu: komenda build i komenda ask, która wczytuje model zapisany w indeksie:

import argparse
import sys
from pathlib import Path

from sentence_transformers import SentenceTransformer

from index import build_index
from search import load_index, search

INDEX_DIR = Path("index")
DOCS_DIR = Path("docs")


def main() -> int:
    parser = argparse.ArgumentParser(description="Wyszukiwarka semantyczna w pamięci")
    sub = parser.add_subparsers(dest="komenda", required=True)
    sub.add_parser("build", help="zbuduj indeks z plików w docs/")
    ask = sub.add_parser("ask", help="zadaj pytanie")
    ask.add_argument("pytanie")
    ask.add_argument("-k", type=int, default=3)
    ask.add_argument("--min-score", type=float, default=0.30)
    args = parser.parse_args()

    if args.komenda == "build":
        print(f"Zbudowano indeks: {build_index(DOCS_DIR, INDEX_DIR)} fragmentów")
        return 0

    index = load_index(INDEX_DIR)
    model = SentenceTransformer(index.model_name)  # ten sam model co przy budowie indeksu
    q = model.encode(args.pytanie, normalize_embeddings=True)
    wyniki = search(index, q, k=args.k, min_score=args.min_score)
    if not wyniki:
        print("Brak wyników powyżej progu podobieństwa.")
        return 1
    for r in wyniki:
        c = r.chunk
        print(f"{r.score:.2f}  {c['source']} › {c['section']}")
        print("      " + c["text"].split("\n\n", 1)[1][:90].replace("\n", " ") + "…")
    return 0


if __name__ == "__main__":
    sys.exit(main())

Uruchomienie:

python cli.py build
# Zbudowano indeks: 19 fragmentów

python cli.py ask "Chcę odzyskać to, co zapłaciłem"
# 0.49  regulamin.md › Rezygnacja i zwrot pieniędzy
#       Zwrot pieniędzy przysługuje w ciągu 14 dni od pierwszej płatności, jeśli nie rozpocząłeś ż…
# 0.43  regulamin.md › Rezygnacja i zwrot pieniędzy
#       Subskrypcję możesz anulować w każdej chwili w ustawieniach konta, a dostęp zachowasz do ko…
# 0.39  regulamin.md › Płatności i faktury
#       Abonament miesięczny kosztuje 49 zł brutto, a roczny 490 zł brutto. Płatność odbywa się ka…

W pytaniu nie ma słów „zwrot”, „pieniędzy” ani „płatności”, a właściwy fragment jest pierwszy. Każdy wynik niesie plik i sekcję, bo zapisaliśmy je jako metadane.

Co zrobić z pytaniami, na które dokumenty nie mają odpowiedzi?

Wyszukiwarka semantyczna zawsze zwróci jakieś najbliższe fragmenty, nawet dla pytania „Jaka jest pogoda w Krakowie?”. Dlatego potrzebny jest próg podobieństwa (--min-score): wyniki poniżej niego uznajemy za brak odpowiedzi. Próg wybieramy z danych. Sprawdziliśmy najlepszy wynik dla dwóch grup pytań:

Pytania spoza tematu nie przekroczyły 0,22, a poprawne trafienia nie zeszły poniżej 0,39. Próg 0,30 leży pomiędzy tymi grupami. Nasze polecenie ask z progiem 0,30 odpowiada dla pogody komunikatem „Brak wyników powyżej progu podobieństwa.” i kodem wyjścia 1.

Próg podobieństwa wyznacz z danych, nie z intuicji. U nas pytania spoza tematu miały najlepszy wynik najwyżej 0,22, a trafne odpowiedzi co najmniej 0,39 (0,41 dla pytań przeformułowanych), więc próg 0,30 leży pośrodku. Z innym modelem lub innymi dokumentami te liczby będą inne: wypisz wyniki dla kilkunastu pytań z obu grup i wybierz środek między nimi.

Dwie uwagi o progu:

• Nie ma uniwersalnej wartości. Skala podobieństwa zależy od modelu, dokumentów i długości fragmentów. Wynik 0,3 z jednym modelem znaczy co innego niż z innym.
• Próg odcina brak tematu, nie złą odpowiedź. Pytanie „Ile trwa kurs od zera?” zwróciło u nas na pierwszym miejscu fragment o zwrocie pieniędzy z wynikiem 0,39, a właściwy fragment (kurs „Python od zera”, 20 godzin) dostał ten sam wynik po zaokrągleniu, 0,39. Oba były ponad progiem, więc wyszukiwarka uznałaby je za odpowiedź. Takie niejednoznaczne pytania łatwo przeoczyć, jeśli nie masz zestawu pytań testowych.

Jak sprawdzić, czy wyszukiwarka działa dobrze?

Tak samo jak w poprzedniej części: zbiór pytań z oczekiwaną odpowiedzią i liczenie trafień. Tym razem w dwóch miarach: czy właściwy fragment jest pierwszy (top-1) i czy jest wśród trzech pierwszych (top-3):

from pathlib import Path

from sentence_transformers import SentenceTransformer

from search import load_index, search

# (pytanie, tekst, który musi być w znalezionym fragmencie)
PYTANIA = [
    ("Czy dostanę zwrot pieniędzy?", "14 dni"),
    ("Jak długo ważny jest link do zmiany hasła?", "ważny przez godzinę"),
    ("Ile kosztuje roczna subskrypcja?", "490 zł"),
    ("Jakie wymagania ma aplikacja na telefonie?", "iOS 15"),
    ("Jak zdobyć certyfikat ukończenia?", "70 procent"),
    # ... w pełnej wersji 10 pytań
]

index = load_index(Path("index"))
model = SentenceTransformer(index.model_name)
top1 = top3 = 0
for pytanie, odpowiedz in PYTANIA:
    wyniki = search(index, model.encode(pytanie, normalize_embeddings=True), k=3)
    teksty = [r.chunk["text"] for r in wyniki]
    top1 += odpowiedz in teksty[0]
    top3 += any(odpowiedz in t for t in teksty)
print(f"top-1: {top1}/{len(PYTANIA)}, top-3: {top3}/{len(PYTANIA)}")

Na naszych 10 pytaniach wyszło top-1: 10/10, top-3: 10/10, a na 5 przeformułowanych pytaniach (np. „Czy mogę uczyć się na iPadzie?” → fragment o systemie iOS) 5/5 w top-1.

Nie ciesz się za wcześnie: pytania pisaliśmy, znając dokumenty, więc wyszły łatwiejsze niż prawdziwe pytania użytkowników. 3 dokumenty i 19 fragmentów to bardzo mały zbiór. Taki wynik mówi, że kod działa i że wybór modelu oraz podziału na fragmenty był rozsądny, a nie że wyszukiwarka sprawdzi się na tysiącu stron.

Wynik 10/10 na pytaniach, które sam napisałeś, niewiele znaczy. Pytania pisaliśmy, znając dokumenty, więc wyszły łatwiejsze niż prawdziwe. Poproś o pytania kogoś, kto dokumentów nie czytał, i zobacz, ile z nich trafia.

Jak przetestować wyszukiwarkę w pytest?

Funkcje split_by_sections i search nie potrzebują modelu, więc testujesz je szybko i bez internetu. W teście search zamiast prawdziwych embeddingów używamy małych, ręcznie wpisanych wektorów, dla których znamy prawidłowy wynik. Podobne podejście (podmiana kosztownego wywołania na kontrolowane dane) opisujemy w artykule o testowaniu kodu z API LLM w pytest.

import numpy as np
import pytest

from chunking import split_by_sections
from search import Index, search

TEKST = """# Dokument

## Zwrot

Zwrot pieniędzy przysługuje w ciągu 14 dni.

## Hasło

Hasło zresetujesz na stronie logowania.
"""


def test_split_by_sections_zachowuje_zrodlo_i_pozycje():
    chunks = split_by_sections(TEKST, "dok.md")

    assert [c.section for c in chunks] == ["Zwrot", "Hasło"]
    assert all(c.source == "dok.md" for c in chunks)
    assert TEKST[chunks[1].start:].startswith("## Hasło")


@pytest.fixture
def index():
    wektory = np.array([[1.0, 0.0], [0.0, 1.0], [0.7071, 0.7071]], dtype=np.float32)
    chunks = [{"id": f"a#{i}", "section": s} for i, s in enumerate(["A", "B", "C"])]
    return Index(wektory, chunks, "fake")


def test_search_zwraca_najlepsze_wyniki_w_kolejnosci(index):
    wyniki = search(index, np.array([1.0, 0.0], dtype=np.float32), k=2)

    assert [r.chunk["section"] for r in wyniki] == ["A", "C"]


def test_search_odrzuca_wyniki_ponizej_progu(index):
    wyniki = search(index, np.array([1.0, 0.0], dtype=np.float32), k=3, min_score=0.9)

    assert [r.chunk["section"] for r in wyniki] == ["A"]


def test_search_zglasza_zly_ksztalt_zapytania(index):
    with pytest.raises(ValueError, match="Zły kształt"):
        search(index, np.array([[1.0, 0.0]], dtype=np.float32))

Uruchom je przez python -m pytest: wtedy bieżący katalog trafia do ścieżki importów i from chunking import ... działa bez dodatkowej konfiguracji. W naszym projekcie przechodzi 6 testów, w tym sprawdzenie, że żaden fragment nie przekracza limitu tokenów i że niezgodny wymiar zapytania kończy się błędem.

Jakie ograniczenia ma wyszukiwarka w pamięci?

Zanim rozbudujesz projekt, warto znać jego granice:

Brute-force (porównanie z każdym wektorem) jest w tym zastosowaniu w pełni wystarczający. Specjalne struktury wyszukiwania w bazach wektorowych mają sens dopiero, gdy dokumentów jest tak dużo, że liczenie wszystkich podobieństw trwa za długo.

Co dalej po tym projekcie?

Masz kompletną wyszukiwarkę, którą możesz rozbudować na kilka sposobów:

• Wynik jako kontekst dla modelu językowego. Najlepsze fragmenty wklejasz do promptu i prosisz model o odpowiedź na podstawie źródeł. To podstawa RAG, o którym piszemy w artykule o bazach wektorowych i RAG w Pythonie.
• API zamiast CLI. Opakuj search w endpoint FastAPI i ładuj model raz przy starcie.
• Lepszy podział i więcej pytań testowych. Dodaj własne dokumenty i pytania zadane przez kogoś, kto ich nie czytał, a potem powtórz pomiar z poprzedniej części.
• Model lokalny zamiast API. Całość działa bez wysyłania danych na zewnątrz, a jeśli chcesz uruchomić również generowanie odpowiedzi lokalnie, zobacz model w Ollamie.

Jeśli chcesz zbudować taki projekt z mentorem i dostać code review swojego kodu, zajrzyj na ścieżkę Python.