vllm-0.10.1
data
moulinette*
+.mypy_cache
+.ruff_cache
--- /dev/null
+install:
+ uv sync
+
+run:
+ uv run rag --help
+
+debug:
+ uv run -m -- pdb -m src
+
+clean:
+ find src -type d \( -name "__pycache__" -or -name ".mypy_cache" -or -name ".ruff_cache" \) -exec rm -r {} +
+
+lint:
+ uv run flake8 src
+ uv run mypy src --warn-return-any --warn-unused-ignores --ignore-missing-imports --disallow-untyped-defs --check-untyped-defs
+
+lint-strict:
+ uv run flake8 src
+ uv run mypy src --strict
+
+check:
+ uv run ty check
+ uv run ruff check
"flake8>=7.3.0",
"mypy>=2.3.1",
"ruff>=0.16.4",
+ "ty>=0.0.75",
"types-tqdm>=4.70.0.20260805",
]
[tool.ruff]
line-length = 79
-lint.select = ["E", "F", "UP", "B", "SIM", "I", "ARG", "N"]
+lint.select = ["E", "F", "UP", "B", "SIM", "I", "ARG", "N", "D"]
+
+[tool.ruff.lint.pydocstyle]
+convention = "google"
[tool.uv]
preview-features = ["format", "check"]
[[tool.mypy.overrides]]
module = ["fire"]
follow_untyped_imports = true
+
+[tool.ty.rules]
+all = "error"
+possibly-missing-import = "ignore"
+"""A retrieval augmented generation CLI."""
+
from rag import main
if __name__ == "__main__":
+"""A retrieval augmented generation CLI."""
+
import logging
import os
import traceback
from rag.evaluate import sources_overlap
from rag.models import (
AnswerSet,
- RagDataset,
+ RagQuestions,
RetrievedQuestion,
SearchAnswers,
SearchResults,
def main() -> None:
+ """Run RAG using python fire, handling all exceptions."""
try:
fire.Fire(RAG())
except Exception as e:
class RAG:
- """A retrieval augmented generation CLI"""
+ """A retrieval augmented generation CLI."""
def __init__(self) -> None:
+ """Set up the caching state."""
self._bm25_store: BM25[Source] | None = None
self._inference_store: InferenceBackend | None = None
index: str = "data/processed",
max_chunk_size: int = 2000,
) -> None:
+ """Index a directory and its contents for later querying."""
if max_chunk_size < 1:
logging.getLogger(__name__).error(
f"Invalid max chunk size {max_chunk_size}"
return self._bm25_store
bm25: BM25[Source] = self._readf(
index + "/index.json",
- lambda s: BM25.from_storage(
+ lambda s: BM25[Source].from_storage(
(key_adapter.validate_json(k), v)
for k, v in tqdm(
storage_adapter.validate_json(s).items(),
def search(
self, query: str, /, k: int = 5, index: str = "data/processed"
) -> None:
+ """Retrieve sources for a querie."""
if k < 1:
logging.getLogger(__name__).error(f"Invalid k {k}")
exit(1)
k: int = 5,
index: str = "data/processed",
) -> None:
+ """Retrieve sources for queries."""
if k < 1:
logging.getLogger(__name__).error(f"Invalid k {k}")
exit(1)
bm25 = self._bm25(index)
- dataset = self._readf(dataset_path, RagDataset.model_validate_json)
+ dataset = self._readf(dataset_path, RagQuestions.model_validate_json)
result = SearchResults(
search_results=[
RetrievedQuestion(
index: str = "data/processed",
max_tokens: int = 100,
) -> None:
+ """Answer a query, retrieving sources and using an LLM."""
if k < 1:
logging.getLogger(__name__).error(f"Invalid k {k}")
exit(1)
/,
max_tokens: int = 100,
) -> None:
+ """Answer queries given search results, using an LLM."""
if max_tokens < 1:
logging.getLogger(__name__).error(
f"Invalid max tokens {max_tokens}"
k: int = 5,
iou: float = 0.05,
) -> None:
+ """Evaluate the student search results for overlap with dataset."""
if k < 1:
logging.getLogger(__name__).error(f"Invalid k {k}")
exit(1)
-"""Useless module where I can shove the required but overcomplicated and
-poorly designed pydantic classes."""
+"""Useless module where I can shove the required poorly pydantic models."""
import uuid
+"""Answering questions using LLMs."""
+
import logging
from collections.abc import Callable
-from typing import Any, cast
+from typing import Any, cast, override
from rag.models import AnsweredQuestion, RetrievedQuestion, Source
class InferenceBackend:
+ """A simple inference backend using transformers."""
+
SYSTEM_PROMPT: str = """You are a codebase and document search assistant.
You answer querries factually using the provided sources.
Answer directly with no reasoning, no <think> tags.
End your answer with a newline."""
def __init__(self, model: str) -> None:
+ """Initialize the backend fetching the given model from huggingface."""
from transformers import AutoModelForCausalLM, AutoTokenizer
self.model: Any = AutoModelForCausalLM.from_pretrained(
model, device_map="auto"
)
- self.tokenizer = AutoTokenizer.from_pretrained(
+ self.tokenizer: Any = AutoTokenizer.from_pretrained(
model, padding_side="left"
)
max_tokens: int,
cb: Callable[[str], None],
) -> str:
- from transformers.generation import BaseStreamer # type: ignore
+ from transformers.generation.streamers import BaseStreamer
class Streamer(BaseStreamer):
def __init__(
self.cb = cb
self.skip = True
+ @override
def put(self, value: Any) -> None:
if self.skip:
self.skip = False
)
self.cb(s.rstrip("\n"))
+ @override
def end(self) -> None:
self.cb("\n")
tokenizer=self.tokenizer,
streamer=Streamer(self, cb),
)
- full_response = self.tokenizer.batch_decode(
- generated_ids,
- skip_special_tokens=True,
- )[0]
+ full_response: str = cast(
+ str,
+ self.tokenizer.batch_decode(
+ generated_ids,
+ skip_special_tokens=True,
+ )[0],
+ )
return full_response[len(prompt) :].strip()
def answer(
max_tokens: int,
cb: Callable[[str], None] = lambda _s: None,
) -> AnsweredQuestion:
+ """Answers the given question."""
return AnsweredQuestion(
question_id=question.question_id,
question=question.question,
+"""Utilities for chunking a file into chunks respecting filetype semantics."""
+
import logging
from enum import Enum, auto
from itertools import count
class FileType(Enum):
+ """A supported file type enum."""
+
PYTHON = auto()
TEXT = auto()
- def line_depth(self, line: str) -> int:
+ def _line_depth(self, line: str) -> int:
match self:
case FileType.PYTHON:
return next(
)
-def chunk_by_depth(
+def _chunk_by_depth(
s: str, max_chunk_size: int, by: FileType, layer: int = 0
) -> list[str]:
chunks = [""]
for line in s.splitlines(True):
- depth = by.line_depth(line)
+ depth = by._line_depth(line)
if depth is None or depth > layer:
chunks[-1] += line
else:
s[i : min(len(chunks), i + max_chunk_size)]
for i in range(0, len(chunks), max_chunk_size)
]
- res = []
+ res: list[str] = []
can_extend = False
for chunk in chunks:
if len(chunk) > max_chunk_size:
- res.extend(chunk_by_depth(chunk, max_chunk_size, by, layer + 1))
+ res.extend(_chunk_by_depth(chunk, max_chunk_size, by, layer + 1))
can_extend = False
elif not can_extend or len(chunk) + len(res[-1]) > max_chunk_size:
res.append(chunk)
def chunk_file(
path: str, max_chunk_size: int, by: FileType
) -> dict[Source, str]:
+ """Appropriately chunk a file of a given filetype into smaller chunks."""
try:
with open(path) as f:
s = f.read()
f"Error while chunking file {path}, skipping ({e})"
)
return {}
- chunks = chunk_by_depth(s, max_chunk_size, by)
+ chunks = _chunk_by_depth(s, max_chunk_size, by)
total_len = 0
res = {}
for chunk in chunks:
+"""Utilities for evaluating search results compared to expected results."""
+
from rag.models import Source
def sources_overlap(a: Source, b: Source, min_iou: float) -> bool:
+ """Check whether two sources overlap with a minimum IoU."""
union_start = min(a.first_character_index, b.first_character_index)
union_end = max(a.last_character_index, b.last_character_index) + 1
union = union_end - union_start
+"""Pydantic models representing the inputs and outputs of the program."""
+
import uuid
+from typing import override
from pydantic import BaseModel, Field
class Source(BaseModel, frozen=True):
+ """A source identifying a file and span within in."""
+
file_path: str
first_character_index: int
last_character_index: int
+ @override
def __str__(self) -> str:
+ """Represent the source in the suggested format."""
return (
f"{self.file_path}"
+ f" [{self.first_character_index}:"
class Question(BaseModel):
+ """A queried question, without further processing."""
+
question_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
question: str
-class RagDataset(BaseModel):
- rag_questions: list[Question]
-
-
class RetrievedQuestion(Question):
+ """A queried question, with search results."""
+
retrieved_sources: list[Source]
class AnsweredQuestion(RetrievedQuestion):
+ """A queried question, with search results and LLM answers."""
+
answer: str
+class RagQuestions(BaseModel):
+ """A set of queried question, without further processing."""
+
+ rag_questions: list[Question]
+
+
class SearchResults(BaseModel):
+ """A set of queried questions, with search results."""
+
search_results: list[RetrievedQuestion]
k: int
class SearchAnswers(BaseModel):
+ """A set of queried questions, with search results and LLM answers."""
+
search_results: list[AnsweredQuestion]
k: int
class DatasetAnswer(BaseModel):
+ """An answer according to the dataset schema."""
+
question_id: str
question: str
answer: str
class AnswerSet(BaseModel):
+ """A set of answers according to the dataset schema."""
+
rag_questions: list[DatasetAnswer]
+"""Retrieval indexing and lookup for a corpus of text."""
+
import math
from collections.abc import Iterable
from dataclasses import dataclass
def words_normalize(s: str) -> Iterable[str]:
+ """Split a string into normalized words."""
start = 0
for i, c in enumerate(s):
if c.isalnum():
def bag_of_words(s: str, stopwords: set[str]) -> BagOfWords:
+ """Split a string into a normalized bag of words, ignoring stopwords."""
res: BagOfWords = {}
for word in words_normalize(s):
if word in stopwords:
def stopwords(lang: str | Iterable[str]) -> set[str]:
+ """Return a usable set of normalized stopwords in lang."""
return {
word
for e in stopwordsiso.stopwords(lang)
@dataclass
class BM25[T]:
+ """A BM25 index for a given corpus allowing for fast lookups."""
+
corpus_words: int
corpus: dict[T, tuple[int, BagOfWords]]
word_usage: dict[str, set[T]]
@staticmethod
def from_storage(storage: Iterable[tuple[T, BagOfWords]]) -> "BM25[T]":
+ """Load an index from its simplifed representation."""
corpus = {}
word_usage: dict[str, set[T]] = {}
corpus_words = 0
return BM25(corpus_words, corpus, word_usage)
def to_storage(self) -> Iterable[tuple[T, BagOfWords]]:
+ """Transform an index to its simplified representation."""
return ((k, v[1]) for k, v in self.corpus.items())
@staticmethod
def from_corpus(
raw_corpus: Iterable[tuple[T, str]], stopwords: set[str] | None = None
) -> "BM25[T]":
- return BM25.from_storage(
+ """Create an index for a given corpus of texts, ignoring stopwrods."""
+ return BM25[T].from_storage(
(k, bow)
for k, s in raw_corpus
if (bow := bag_of_words(s, stopwords if stopwords else set()))
)
- def idf(self, q: str) -> float:
+ def _idf(self, q: str) -> float:
n_q = len(self.word_usage.get(q, set()))
n = len(self.corpus)
quotient = (n - n_q + 0.5) / (n_q + 0.5)
return math.log(quotient + 1)
def word_score(self, word: str, ident: T) -> float:
+ """Return the score of a document for a single query word."""
doc_words, doc = self.corpus[ident]
f_q = doc.get(word, 0)
freq = f_q * (self.k + 1)
avgdl = self.corpus_words / len(self.corpus)
freq_bias = 1 - self.b + self.b * (doc_words / avgdl)
quotient = freq / (f_q + self.k * freq_bias)
- return self.idf(word) * quotient
+ return self._idf(word) * quotient
- def score(self, querry: str, doc: T) -> float:
- return sum(self.word_score(q, doc) for q in words_normalize(querry))
+ def score(self, query: str, doc: T) -> float:
+ """Return the score of a document for a query."""
+ return sum(self.word_score(q, doc) for q in words_normalize(query))
def word_scores(self, word: str) -> dict[T, float]:
+ """Return all the non-zero scoring documents and theirs scores."""
res = {}
k_plus_1 = self.k + 1
avgdl = self.corpus_words / len(self.corpus)
freq_bias_add = 1 - self.b
freq_bias_mul = self.b / avgdl
- idf = self.idf(word)
+ idf = self._idf(word)
for ident in self.word_usage.get(word, set()):
doc_words, doc = self.corpus[ident]
f_q = doc[word]
res[ident] = idf * quotient
return res
- def scores(self, querry: str) -> dict[T, float]:
- # Old slow impl:
- # return {k: self.score(querry, k) for k in self.corpus}
+ def scores(self, query: str) -> dict[T, float]:
+ """Return the non-zero-scoring documents and their score."""
res: dict[T, float] = {}
- for word in words_normalize(querry):
+ for word in words_normalize(query):
for k, v in self.word_scores(word).items():
res[k] = res.get(k, 0.0) + v
return res
- def best_k(self, querry: str, k: int) -> list[T]:
- k = min(len(self.corpus), k)
+ def best_k(self, query: str, k: int) -> list[T]:
+ """Return the k best non-zero scoring documents in the corpus."""
scores = sorted(
- self.scores(querry).items(),
+ self.scores(query).items(),
key=lambda e: e[1],
reverse=True,
)[:k]
{ name = "flake8" },
{ name = "mypy" },
{ name = "ruff" },
+ { name = "ty" },
{ name = "types-tqdm" },
]
{ name = "flake8", specifier = ">=7.3.0" },
{ name = "mypy", specifier = ">=2.3.1" },
{ name = "ruff", specifier = ">=0.16.4" },
+ { name = "ty", specifier = ">=0.0.75" },
{ name = "types-tqdm", specifier = ">=4.70.0.20260805" },
]
{ url = "https://files.pythonhosted.org/packages/41/c4/a12e1d9b387fb0c40a57116db82b457e8c771cb419163cda29204d74a595/transformers-5.15.1-py3-none-any.whl", hash = "sha256:b7cdf238ff583e3a58dbc7fa34da1aaf091ce063141f65a30538160bd5afe93f", size = 11749582, upload-time = "2026-08-19T11:28:16.726Z" },
]
+[[package]]
+name = "ty"
+version = "0.0.75"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/81/d0/d0c96f898d6974a4a3569ab3efdf9512c04ad99f9203effb55f72497fe97/ty-0.0.75.tar.gz", hash = "sha256:4c5eead33dfbf6e2ebb4f400f74b51ffc9bab702a6f23ddb648a1cbb740387e3", size = 6868326, upload-time = "2026-08-26T20:23:40.399Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/cb/6c/b12d03505f17581f0cfa3c12273fe34c1d67b36dfda1bc561a6bdc16512b/ty-0.0.75-py3-none-linux_armv6l.whl", hash = "sha256:e5409f50db2246fd4bd039d93d261e0cfa1daa554a4fb77256f91072c570349a", size = 12972606, upload-time = "2026-08-26T20:22:59.716Z" },
+ { url = "https://files.pythonhosted.org/packages/d1/aa/30f11eecd9215a9f87e8fe8baaf48f3ce905f5d75b8e4aac70f0091f130c/ty-0.0.75-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:5e7b8b3472fb9bb2eeab314984b265df08a7a9d518867a9e6020eebc06570be2", size = 12527158, upload-time = "2026-08-26T20:23:02.767Z" },
+ { url = "https://files.pythonhosted.org/packages/f2/11/7fd7001b0b5c6610bfbad7357e47d5fe6f82d4e84e94c53776a478f5e9f8/ty-0.0.75-py3-none-macosx_11_0_arm64.whl", hash = "sha256:c6ccf34169821fe0d23e3360deeef981d217963412f1d087b9bdd32ec57f7a57", size = 12400533, upload-time = "2026-08-26T20:23:04.965Z" },
+ { url = "https://files.pythonhosted.org/packages/fd/7f/1e284ea3d348d7be02f12d83bc22ed9ef193033f863f05b64db99027f141/ty-0.0.75-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:842ebb41e9c6c334b40768704e20b1a69d5c6b08805b289d5e0e2565f49f2de1", size = 12420592, upload-time = "2026-08-26T20:23:07.427Z" },
+ { url = "https://files.pythonhosted.org/packages/2d/ab/d813271543370c47fd74b5118f2066ab32b0983e907b1821f3f9a6d0fa7f/ty-0.0.75-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:cf7a5a723c5f1e0fab4ffbfe9bd95123a526ed48f206e5f25cb2161ca294007a", size = 12739219, upload-time = "2026-08-26T20:23:09.809Z" },
+ { url = "https://files.pythonhosted.org/packages/31/5b/95b49cc5570fd92a7bf63732f649b31906158721e03c7fcb1b5be74ee3bf/ty-0.0.75-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:54382f98e5da292fcd7104391afef5105c35bb2f312e29bea6f5fa419935255c", size = 13494046, upload-time = "2026-08-26T20:23:12.191Z" },
+ { url = "https://files.pythonhosted.org/packages/2c/0d/502d2dd68173cf020e1ad2bdbab9544c86776de0b0e2ed15f8c2fe006e3d/ty-0.0.75-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:ac13b180dc2aade2cd243f56b01650e78bf091a2e522ad3bc947245d7837c613", size = 13938899, upload-time = "2026-08-26T20:23:14.764Z" },
+ { url = "https://files.pythonhosted.org/packages/20/5b/f3b12a25c07224456219fc2bd20db0ad7e40b304be0ff6aad728da0135f9/ty-0.0.75-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:752df7951a443219d7f1ff817e3723c85d428565ff449e08a7a93ba821661526", size = 13656711, upload-time = "2026-08-26T20:23:17.145Z" },
+ { url = "https://files.pythonhosted.org/packages/51/7b/f090ad306e2b15a07b332d647138c5264b89d9758855ecce8b8a10bcb153/ty-0.0.75-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1fd399feedf7cee816563c1baec45fc1c0b3c89f1ea42364920b688004b5b7da", size = 13093499, upload-time = "2026-08-26T20:23:19.489Z" },
+ { url = "https://files.pythonhosted.org/packages/f1/4b/f69b99aaaca0c7c65d5f114b186b26b21666f767b0c69eec99a2bdccc061/ty-0.0.75-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:d7625f6f56c7dc1e873579fdc9e432a0e21e302afe847ab60704d2303442a92e", size = 13520580, upload-time = "2026-08-26T20:23:21.789Z" },
+ { url = "https://files.pythonhosted.org/packages/b6/e7/692c5f905c0345a15d2255fc74066d660f030254ae8dcdaf33f5a5c2f279/ty-0.0.75-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:89e7d527e95a2534b70cae29e94c104b84082760ea05927d23bb87280969c104", size = 12524095, upload-time = "2026-08-26T20:23:24.026Z" },
+ { url = "https://files.pythonhosted.org/packages/ba/9a/f42b12cf265ea95344bf554764c4791cfb273bdd628aadd7c209af7cadc3/ty-0.0.75-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:0843f134440740706e01bee5f88f4cfc10e9b018bddb9e4ef4c12dc9fc0c9aef", size = 12756591, upload-time = "2026-08-26T20:23:26.126Z" },
+ { url = "https://files.pythonhosted.org/packages/7c/5e/9b180c133cb9cce48179a7d2bf9e1802d992aa8176a918e0e05205760b42/ty-0.0.75-py3-none-musllinux_1_2_i686.whl", hash = "sha256:1bd0ec0e50ee1376875c88891efe6f549c3560fa5b2ddad79a425cd5a6218b9c", size = 12998754, upload-time = "2026-08-26T20:23:28.353Z" },
+ { url = "https://files.pythonhosted.org/packages/39/f6/3c6ef5dd550103e29905121c67fb96a374564f31a2f44c6faa1af98c2d61/ty-0.0.75-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:1f9eafd561f90110d5e29f589ec3e956c4686e2f6631348d99276436f5cbe4d1", size = 13316474, upload-time = "2026-08-26T20:23:30.857Z" },
+ { url = "https://files.pythonhosted.org/packages/bb/52/12776337874c821076bd5368e352ccd9e67174790abe3b856f749cb3524b/ty-0.0.75-py3-none-win32.whl", hash = "sha256:05063a6fafe2154b794a7f964515d148e51acd186d72d4a3acd347ee9fa19336", size = 12316315, upload-time = "2026-08-26T20:23:33.528Z" },
+ { url = "https://files.pythonhosted.org/packages/53/e6/bb51e16af5c7138c9f52f8f3d0a401a371c6798d092e3b74926f186a9814/ty-0.0.75-py3-none-win_amd64.whl", hash = "sha256:81cf1ba5f6b7536ad56747865214255d9bc8e80533a689dbb9ddeaad464b09f1", size = 12917267, upload-time = "2026-08-26T20:23:35.978Z" },
+ { url = "https://files.pythonhosted.org/packages/39/73/4542f829107468b5de4231af67f29927c093bfad11f3c1e5b2c08fb1206b/ty-0.0.75-py3-none-win_arm64.whl", hash = "sha256:541c9af5b7a0ad23d15ec315a7da81150833c359f48124ed3789ff25eacd6f42", size = 12711024, upload-time = "2026-08-26T20:23:38.159Z" },
+]
+
[[package]]
name = "typer"
version = "0.27.1"
projects / axy / ft / rag.git / commitdiff