Spaces:
Sleeping
Sleeping
File size: 4,019 Bytes
2bb4a0d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | """ARIA RAG engine.
Loads markdown knowledge files, chunks them by `##` heading (then by
~300-token windows with overlap), embeds with all-MiniLM-L6-v2, and serves
cosine-similarity search via a FAISS inner-product index over normalized
vectors. The index is built once at startup — the knowledge base is small,
so this takes seconds.
"""
from __future__ import annotations
import os
import re
from dataclasses import dataclass
KNOWLEDGE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "knowledge")
EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
CHUNK_WORDS = 220 # ~300 tokens
OVERLAP_WORDS = 40 # ~50 tokens
TOP_K = 4
@dataclass
class Chunk:
source: str # filename, e.g. "genie.md"
heading: str # nearest ## heading ("" for file preamble)
text: str # chunk text with heading prepended
def _split_sections(markdown: str) -> list[tuple[str, str]]:
"""Split a markdown document into (heading, body) pairs on `##` headings."""
parts = re.split(r"(?m)^##\s+", markdown)
sections: list[tuple[str, str]] = []
preamble = parts[0].strip()
if preamble:
# Strip the top-level `# Title` line but keep any intro text.
preamble_body = re.sub(r"(?m)^#\s+.*$", "", preamble).strip()
title = ""
m = re.search(r"(?m)^#\s+(.*)$", preamble)
if m:
title = m.group(1).strip()
if preamble_body:
sections.append((title, preamble_body))
for part in parts[1:]:
lines = part.splitlines()
heading = lines[0].strip() if lines else ""
body = "\n".join(lines[1:]).strip()
if body:
sections.append((heading, body))
return sections
def _windows(words: list[str], size: int, overlap: int):
"""Yield overlapping word windows of `size` with `overlap` words shared."""
step = max(size - overlap, 1)
start = 0
while True:
yield words[start:start + size]
if start + size >= len(words):
break
start += step
def load_chunks() -> list[Chunk]:
"""Read every .md file in knowledge/ and return its chunks."""
chunks: list[Chunk] = []
for name in sorted(os.listdir(KNOWLEDGE_DIR)):
if not name.endswith(".md"):
continue
with open(os.path.join(KNOWLEDGE_DIR, name), encoding="utf-8") as f:
text = f.read()
for heading, body in _split_sections(text):
words = body.split()
for piece in _windows(words, CHUNK_WORDS, OVERLAP_WORDS):
if not piece:
continue
chunk_text = " ".join(piece)
if heading:
chunk_text = f"{heading}\n{chunk_text}"
chunks.append(Chunk(source=name, heading=heading, text=chunk_text))
return chunks
class RagIndex:
"""In-memory FAISS index over the knowledge chunks (cosine similarity)."""
def __init__(self) -> None:
# Heavy imports kept local so chunking can be unit-tested without them.
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
self._np = np
self.chunks = load_chunks()
if not self.chunks:
raise RuntimeError(f"No knowledge chunks found in {KNOWLEDGE_DIR}")
self.model = SentenceTransformer(EMBED_MODEL, device="cpu")
vectors = self.model.encode(
[c.text for c in self.chunks],
normalize_embeddings=True,
show_progress_bar=False,
)
vectors = np.asarray(vectors, dtype="float32")
self.index = faiss.IndexFlatIP(vectors.shape[1]) # cosine via normalized IP
self.index.add(vectors)
def search(self, query: str, k: int = TOP_K) -> list[Chunk]:
q = self.model.encode([query], normalize_embeddings=True)
q = self._np.asarray(q, dtype="float32")
_, idx = self.index.search(q, min(k, len(self.chunks)))
return [self.chunks[i] for i in idx[0] if i != -1]
|