プロジェクト

ローカルのエンベディングで作る、テキストフォルダ上の小さなRAGシステム

ノートのフォルダをチャンク化し、CPU上のall-MiniLM-L6-v2でチャンクを埋め込み、素のnumpyによるコサイン索引で検索し、上位チャンクをOllama経由でローカルモデルに渡して、出典付きの回答を得ます。

実行環境ノートPCのCPU、5分以内(生成ステップにはOllamaが必要)
所要時間約3時間
使用ライブラリpython, numpy, huggingface
追加パッケージsentence-transformers、httpx

作るもの

検索拡張生成(RAG)は、1つのふりをした2つのプログラムです。1つ目は文書の山をベクトルに変え、質問に最も似た少数の断片を見つけます。2つ目はその断片から回答を書きます。このプロジェクトでは、ベクトルデータベースもフレームワークもホスト型APIも使わずに、本物らしく振る舞う最小のサイズで両方を作ります。検索側は22 Mパラメータのエンベディングモデルで、すべてCPU上で動きます。生成側はOllamaの下でローカルに動くモデルと通信します。

コーパスは互いに無関係な話題(サワードウ、自転車のギア、灯台、堆肥、潮汐など)の短いノート9本で、最初のブロックがcorpus/フォルダに書き出します。無関係な話題であることが肝心です。3つのテストクエリがそれぞれちょうど1つのファイルからチャンクを引き出せば、目を凝らさなくても検索が機能していると分かります。その後フォルダを自分のノートに差し替えても、他には何も変わりません。

エンベディングのレッスンとLLMのレッスンに慣れている必要があります。RAGパイプラインのチェックリストには、このプロジェクトが代わりに決めてしまっている判断が列挙されているので、コーパスが大きくなったときに何を見直すべきかが分かります。

必要なもの

  • Python 3.10以降とnumpy、sentence-transformers(pip install sentence-transformersで、torchとtransformersも入ります)。初回使用時にエンベディングモデル(約90 MB)をダウンロードします。以後はローカルキャッシュからオフラインで動きます。
  • 生成ステップ用のhttpx(pip install httpx)。
  • 生成のみ: Ollamaをインストールして起動し、モデルをpullしておきます。当サイトではqwen2.5:7b(Qwen2.5 7B Instruct、7.6 Bパラメータ、Ollamaビルドでは32 kコンテキスト、Apache-2.0、2024年9月公開)を使いました。デフォルトの4ビット量子化で4.7 GBのダウンロードで、約6 GBの空きメモリが必要です。Ollamaが提供する指示調整済みモデルなら何でも動きます。別のモデルを使うにはOLLAMA_MODELを設定してください。
  • GPUは不要です。エンベディングモデルはこのコーパスに対してCPUで1秒をはるかに下回る時間で動きます。OllamaはApple SiliconではGPUを自動的に使い、それ以外ではCPUにフォールバックします。

検索ブロックは、Apple SiliconのMac mini上で、Field Guideのパイプライン環境を使い、空のディレクトリから実行しました。生成ブロックは同じマシンでローカルのOllamaに対して実行しました。所要時間は各ブロックに記載しています。

設計

corpus/9 text filesチャンク化2文、ストライド1MiniLM-L6-v2384次元、単位長index.npy36 x 384 float32クエリ時質問MiniLM-L6-v2同じモデルindex @ qargsort, top kOllamaprompt = sources + Q[n]の出典付き回答
索引は一度だけ作り、クエリは何度でも。生成モデルが必要なのは網掛けのステップだけ。

チャンク化。 各ファイルを文に分割し、連続する2文のウィンドウをストライド1で切り出してチャンクとします。5文のファイルからは4チャンクができ、すべての文が2つのチャンクに現れます。この重なりにより、文の境界をまたいで分かれた事実も少なくとも1つのチャンクの中に収まります。これほど短いノートにはこれが適切なサイズです。長い文書では段落ごと、あるいは200から500トークンの予算でチャンク化することになり、そのトレードオフはチェックリストのページで扱っています。

エンベディング。 sentence-transformers/all-MiniLM-L6-v2(22.7 Mパラメータ、256ワードピースのコンテキスト、Apache-2.0、2021年公開)は、チャンクを384次元のベクトルに写像します。normalize_embeddings=Trueによりすべてのベクトルが単位長になるので、2つのベクトルの内積がそのままコサイン類似度になります。これで検索全体が行列ベクトル積index @ q1回と、それに続くargsortになります。36チャンクなら1万4千回の乗算です。これは数十万チャンクまで正しい設計であり続け、ほとんどの個人のコーパスはそこまで到達しません。それを超えたらFAISSのような近似索引やベクトルデータベースを追加しますが、その周りのコードは変わりません。

なぜクエリとチャンクに同じモデルを使うのか。 ベクトル空間は1つのモデルに対してしか意味を持ちません。別のモデル、あるいは同じモデルの別バージョンでクエリを埋め込むと、索引と比較できない数値になります。モデル名を索引と一緒に保存してください。

生成。 プロンプトは検索されたチャンクに番号を振り、ファイル名を示し、それらだけから回答して番号で出典を引用するようモデルに指示します。温度0により出力は再現可能になります。それでもモデルの回答はモデルの回答です。検索が保証するのは、モデルが正しいテキストを見たことであって、そこから正しく推論したことではありません。3つのテスト質問のうち1つがまさにその失敗を示しており、注意深く読む価値があります。

コード

コーパス、索引、検索器を作る

このブロックは9つのファイルを書き出し、チャンク化し、索引を作り、3つのクエリに対して上位3チャンクをスコアとファイル名付きで出力します。生成ステップのためにindex.npyとchunks.tsvを保存します。

# rag.py: write a small corpus, chunk it, embed it locally, retrieve by cosine similarity.
import re
import time
from pathlib import Path

import numpy as np
from sentence_transformers import SentenceTransformer

SEED = 0
np.random.seed(SEED)

# --- 1. A corpus of nine short notes on unrelated subjects ---
CORPUS = {
    "sourdough.txt": (
        "A sourdough starter is a culture of wild yeast and lactic acid bacteria kept alive in flour and water. "
        "Feed it once a day at room temperature or once a week in the fridge. "
        "A starter is ready to bake with when it doubles in volume within four to six hours of feeding. "
        "Bulk fermentation at 24 degrees Celsius takes about five hours for a dough with 20 percent starter. "
        "Bake in a preheated cast iron pot at 240 degrees Celsius, lid on for 20 minutes and lid off for 20 more."
    ),
    "bicycle-gears.txt": (
        "A bicycle drivetrain multiplies pedal cadence by the ratio of chainring teeth to sprocket teeth. "
        "A 50 tooth chainring with an 11 tooth sprocket gives a ratio of 4.5, which is a fast gear for descents. "
        "A 34 tooth chainring with a 34 tooth sprocket gives a ratio of 1.0 for steep climbs. "
        "Cross-chaining, running the big chainring with the big sprocket, bends the chain and wears it faster. "
        "A chain stretched beyond 0.75 percent of its nominal length should be replaced before it damages the cassette."
    ),
    "lighthouse.txt": (
        "The Eddystone Rocks lie 14 kilometres south of Plymouth and have carried four successive lighthouses. "
        "Smeaton's tower of 1759 was built from interlocking granite blocks shaped like tree trunks. "
        "It stood until 1877, when erosion of the rock beneath it forced a replacement. "
        "The current tower, designed by James Douglass, was lit in 1882 and automated in 1982. "
        "Smeaton's upper section was rebuilt on Plymouth Hoe, where it is open to visitors."
    ),
    "compost.txt": (
        "A compost heap needs a mix of nitrogen-rich green material and carbon-rich brown material. "
        "A ratio of roughly one part green to two parts brown by volume keeps the heap from going slimy. "
        "Turn the heap every two weeks to let air in, which keeps the aerobic bacteria working. "
        "A hot heap reaches 60 degrees Celsius in its core, which kills most weed seeds. "
        "Finished compost is dark, crumbly and smells of forest floor, usually after six to twelve months."
    ),
    "tides.txt": (
        "Tides are caused by the gravitational pull of the Moon and, to a lesser extent, the Sun. "
        "Spring tides, with the largest range, occur at new and full moon when the two pulls align. "
        "Neap tides, with the smallest range, occur at the quarter moons. "
        "The Bristol Channel has one of the largest tidal ranges in the world, up to 15 metres at Avonmouth. "
        "Tide tables predict high water to within a few minutes years in advance using harmonic analysis."
    ),
    "espresso.txt": (
        "A standard double espresso uses 18 grams of ground coffee to produce 36 grams of liquid in 25 to 30 seconds. "
        "Grind finer to slow the shot down and coarser to speed it up. "
        "Water should be between 92 and 96 degrees Celsius at the group head. "
        "A shot that runs too fast tastes sour and thin, while one that runs too slow tastes bitter and harsh. "
        "Clean the portafilter basket and backflush the machine weekly to keep old oils from spoiling the taste."
    ),
    "knots.txt": (
        "The bowline makes a fixed loop that will not slip and is easy to untie after loading. "
        "The clove hitch secures a rope to a post quickly but can work loose under a varying load. "
        "A figure of eight is the standard stopper knot and is also tied on a bight for climbing. "
        "The sheet bend joins two ropes of different diameters more reliably than a reef knot. "
        "A knot can reduce the breaking strength of a rope by 30 to 50 percent, so leave a margin."
    ),
    "honeybees.txt": (
        "A honeybee colony in summer holds one queen, up to 50,000 workers and a few hundred drones. "
        "Workers forage within about three kilometres of the hive and communicate flower locations with a waggle dance. "
        "The queen lays up to 2,000 eggs a day at the height of the season. "
        "A colony prepares to swarm when it runs out of space, raising a new queen in peanut-shaped cells. "
        "Beekeepers inspect hives every seven to ten days in spring to catch swarm preparations early."
    ),
    "woodstove.txt": (
        "Firewood should be seasoned to below 20 percent moisture before burning, which takes one to two years for oak. "
        "A moisture meter pressed into a freshly split face gives the most honest reading. "
        "A top-down fire, with kindling above the logs, lights cleanly with less smoke. "
        "Close the air control gradually once the fire is established to lengthen the burn. "
        "Have the chimney swept at least once a year to remove creosote, which is a fire risk."
    ),
}
corpus_dir = Path("corpus")
corpus_dir.mkdir(exist_ok=True)
for name, text in CORPUS.items():
    (corpus_dir / name).write_text(text)
print(f"wrote {len(CORPUS)} files to {corpus_dir}/")


# --- 2. Chunk each file into overlapping windows of sentences ---
def chunk_text(text: str, window: int = 2, stride: int = 1) -> list[str]:
    sentences = [s.strip() for s in re.split(r"(?<=[.!?])\s+", text) if s.strip()]
    chunks = []
    for start in range(0, max(1, len(sentences) - window + 1), stride):
        chunks.append(" ".join(sentences[start:start + window]))
    return chunks


chunks, sources = [], []
for path in sorted(corpus_dir.glob("*.txt")):
    for c in chunk_text(path.read_text()):
        chunks.append(c)
        sources.append(path.name)
print(f"{len(chunks)} chunks from {len(set(sources))} files")

# --- 3. Embed every chunk once and L2-normalise, so dot product = cosine ---
t0 = time.perf_counter()
model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", device="cpu")
index = model.encode(chunks, batch_size=32, normalize_embeddings=True, show_progress_bar=False)
index = np.asarray(index, dtype=np.float32)
print(f"index shape {index.shape}, built in {time.perf_counter() - t0:.1f} s")
np.save("index.npy", index)
Path("chunks.tsv").write_text("\n".join(f"{s}\t{c}" for s, c in zip(sources, chunks)))


# --- 4. Retrieve: embed the query, score every chunk, take the top k ---
def retrieve(query: str, k: int = 3) -> list[tuple[float, str, str]]:
    q = model.encode([query], normalize_embeddings=True)[0]
    scores = index @ q
    top = np.argsort(-scores)[:k]
    return [(float(scores[i]), sources[i], chunks[i]) for i in top]


QUERIES = [
    "How hot does the oven need to be for sourdough?",
    "When should I replace a worn bicycle chain?",
    "Why does the Bristol Channel have such big tides?",
]
for query in QUERIES:
    print(f"\nQ: {query}")
    for score, source, chunk in retrieve(query):
        print(f"  {score:.3f}  [{source}]  {chunk[:90]}...")

期待される出力です。Mac miniで4.0秒、そのうちモデルの読み込みと36チャンクの埋め込みは0.3秒です。

wrote 9 files to corpus/
36 chunks from 9 files
index shape (36, 384), built in 0.3 s

Q: How hot does the oven need to be for sourdough?
  0.547  [sourdough.txt]  Bulk fermentation at 24 degrees Celsius takes about five hours for a dough with 20 percent...
  0.466  [sourdough.txt]  A starter is ready to bake with when it doubles in volume within four to six hours of feed...
  0.416  [sourdough.txt]  A sourdough starter is a culture of wild yeast and lactic acid bacteria kept alive in flou...

Q: When should I replace a worn bicycle chain?
  0.455  [bicycle-gears.txt]  Cross-chaining, running the big chainring with the big sprocket, bends the chain and wears...
  0.320  [bicycle-gears.txt]  A 34 tooth chainring with a 34 tooth sprocket gives a ratio of 1.0 for steep climbs. Cross...
  0.307  [bicycle-gears.txt]  A bicycle drivetrain multiplies pedal cadence by the ratio of chainring teeth to sprocket ...

Q: Why does the Bristol Channel have such big tides?
  0.749  [tides.txt]  The Bristol Channel has one of the largest tidal ranges in the world, up to 15 metres at A...
  0.677  [tides.txt]  Neap tides, with the smallest range, occur at the quarter moons. The Bristol Channel has o...
  0.510  [tides.txt]  Tides are caused by the gravitational pull of the Moon and, to a lesser extent, the Sun. S...

すべてのヒットが正しいファイルから来ています。ただしスコアを見てください。サワードウの上位チャンクは0.547、チェーンの上位チャンクは0.455で、どちらのチャンクも確かに答えを含んでいますが、それは2文目にあり、90文字のプレビューでは切れてしまっています。このモデルのコサインスコアは確率ではなく、完全な一致でも0.8を超えることはまれです。自分のデータで測らずに、きりのよい数値で閾値を切らないでください。チャンク化の重なりのおかげで、類似度の大半を担った発酵の文と並んで、240度の文が上位チャンクに入りました。

出典付きの回答を生成する

このブロックには起動中のOllamaが必要です。保存した索引を再読み込みし、各質問について上位4チャンクを検索し、番号付きのプロンプトを組み立てて、Ollamaの/api/generateルートにPOSTします。サーバーのアドレスとモデル名は環境変数から読み、適切なデフォルトを持ちます。

# needs-api
# answer.py: retrieve with the saved index, then ask a local model through Ollama to answer with citations.
import os
from pathlib import Path

import httpx
import numpy as np
from sentence_transformers import SentenceTransformer

OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://127.0.0.1:11434")
OLLAMA_MODEL = os.environ.get("OLLAMA_MODEL", "qwen2.5:7b")
TOP_K = 4

# Load what rag.py saved: the normalised chunk embeddings and the chunk texts with their file names.
index = np.load("index.npy")
rows = [line.split("\t", 1) for line in Path("chunks.tsv").read_text().splitlines()]
sources, chunks = zip(*rows)
embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", device="cpu")


def retrieve(query: str, k: int = TOP_K) -> list[tuple[str, str]]:
    q = embedder.encode([query], normalize_embeddings=True)[0]
    top = np.argsort(-(index @ q))[:k]
    return [(sources[i], chunks[i]) for i in top]


def build_prompt(question: str, hits: list[tuple[str, str]]) -> str:
    context = "\n\n".join(f"[{i + 1}] ({src})\n{text}" for i, (src, text) in enumerate(hits))
    return (
        "Answer the question using only the numbered sources below. "
        "Cite the sources you used as [1], [2] and so on after the sentence they support. "
        "If the sources do not contain the answer, say so.\n\n"
        f"Sources:\n{context}\n\nQuestion: {question}\nAnswer:"
    )


def ask_ollama(prompt: str) -> str:
    response = httpx.post(
        f"{OLLAMA_URL}/api/generate",
        json={"model": OLLAMA_MODEL, "prompt": prompt, "stream": False,
              "options": {"temperature": 0, "num_predict": 200}},
        timeout=120.0,
    )
    response.raise_for_status()
    return response.json()["response"].strip()


QUESTIONS = [
    "How hot does the oven need to be for sourdough?",
    "When should I replace a worn bicycle chain?",
    "Why does the Bristol Channel have such big tides?",
]
for question in QUESTIONS:
    hits = retrieve(question)
    answer = ask_ollama(build_prompt(question, hits))
    print(f"\nQ: {question}\nA: {answer}")
    print("sources: " + ", ".join(f"[{i + 1}] {src}" for i, (src, _) in enumerate(hits)))

実行

pip install numpy sentence-transformers httpx
python rag.py

次に、生成ステップのために別のターミナルで実行します。

ollama pull qwen2.5:7b
ollama serve

ollama serveは、Ollamaのデスクトップアプリケーションがまだ起動していない場合にだけ必要です。最初のターミナルに戻って実行します。

python answer.py

別のモデル、たとえばメモリ8 GBのマシンで小さめのモデルを使うには、実行前にOLLAMA_MODEL=qwen2.5:3bを設定します。

テスト

rag.pyの出力は上に示したとおりで、所要時間を除いて1行ずつ一致するはずです。answer.pyについては、qwen2.5:7bがMac miniで生成したものが以下です。3つの質問で合計22秒かかり、ほぼすべてが生成でした。

Q: How hot does the oven need to be for sourdough?
A: According to the sources provided, the oven needs to be preheated to 240 degrees Celsius for baking the sourdough. This information is cited from [1].
sources: [1] sourdough.txt, [2] sourdough.txt, [3] sourdough.txt, [4] sourdough.txt

Q: When should I replace a worn bicycle chain?
A: You should replace a worn bicycle chain when it stretches beyond 0.75 percent of its nominal length, as this can lead to damage of the cassette before it becomes fully worn out [1].
sources: [1] bicycle-gears.txt, [2] bicycle-gears.txt, [3] bicycle-gears.txt, [4] bicycle-gears.txt

Q: Why does the Bristol Channel have such big tides?
A: The Bristol Channel has such big tides because it experiences spring tides, which have the largest range, occurring at new and full moon when the gravitational pulls of the Moon and the Sun align [3, 4].
sources: [1] tides.txt, [2] tides.txt, [3] tides.txt, [4] tides.txt

最初の2つの回答は正しく、正しい出典を引用しており、質問した人は誰でも納得するでしょう。学ぶべきは3つ目です。コーパスはブリストル海峡の潮差が大きいとは言っています。なぜかは一度も言っていません(本当の理由は海峡の漏斗状の形と潮汐周期との共振で、どのファイルにもありません)。指示は、出典に答えが含まれていなければそう認めるようにというものでした。ところがモデルは、手近にある最も近い仕組みである大潮を取り上げ、もっともらしく見える引用を添えて説明として提示しました。検索は完璧に仕事をしました。生成は正しい出典から自信に満ちた誤答を作り出しました。これがこの分野でいうハルシネーションの定義であり、引用はそれを防ぎませんでした。あなたが作るどのRAGシステムも、ときにはこれをやります。これを捕まえるチェックは、コーパスに答えのない質問からなる評価セットと、引用をたどって読む人です。

別のモデル、あるいは別の日のOllamaビルド上の同じモデルは、3つの回答すべてを異なる言い回しにするかもしれません。比べるべきは引用と、失敗の形です。

試してみること

  • corpus_dirを自分のノートのフォルダか、プロジェクトのdocs/ディレクトリに向けてください。チャンク化を段落単位(text.split("\n\n"))に切り替え、チャンク数と索引の構築時間を出力してください。数千チャンクの索引でも、CPUで1分以内に構築できます。
  • コーパスに答えのない質問(「水銀の沸点は何度か」)を追加し、モデルがそれを認めるかを確認してください。認めるまでプロンプトを締め、最初の2つの回答が悪化していないことを確認してください。
  • 各クエリについて4番目のヒットのスコアを出力してください。それが上位ヒットと別のファイルから来ている場合、プロンプトにノイズを加えています。固定のkの代わりにTOP_K = 2と、上位スコアに対する相対的なスコアの足切り(たとえばその80%)を試してください。
  • コサイン検索をBM25のキーワードスコアリング(rank_bm25パッケージ)に置き換え、チェーンの質問について上位3件を比べてください。次に2つを順位で組み合わせてください。ハイブリッド検索は通常の次のステップで、チェックリストのページがその理由を説明しています。
  • モデル名をヘッダー行としてchunks.tsvに保存し、読み込んだエンベッダーと異なる場合はanswer.pyの実行を拒否するようにしてください。エンベディングモデルを更新したときの、気づかない不一致を防げます。
  • プロンプトはそのままに、Ollamaの呼び出しをホスト型APIに置き換えて、3つの回答を比べてください。システムの検索側はまったく変わりません。

まとめ

  • RAGシステムは検索器と生成器です。別々に作り、別々にテストします。
  • 重なりのある文ウィンドウ、sentence-transformersのモデル、numpyの行列積1回で、数十万チャンクまでのコーパスに対する完全な検索器になります。
  • エンベディングを一度正規化してコサイン類似度を内積にし、クエリは常に索引と同じモデルで埋め込みます。
  • プロンプトでは検索したチャンクに番号を振り、番号による引用を求め、その引用を自分で確認します。
  • 正しい検索は正しい回答を保証しません。コーパスに答えのない質問をテストセットに残しておきます。