import fire
import pydantic
+from pydantic_core import ValidationError
-from rag.chunking import FileChunkID, FileType, chunk_file
-from rag.retrieval import BM25, BagOfWords
+from rag.models import (
+ AnsweredQuestion,
+ MinimalSource,
+ RagDataset,
+ RetrievedQuestion,
+ SearchResults,
+)
+from rag.chunking import FileType, chunk_file
+from rag.retrieval import BM25, BagOfWords, fetch_stopwords
def main() -> None:
- fire.Fire(RAG)
+ fire.Fire(RAG())
storage_adapter = pydantic.TypeAdapter(dict[bytes, BagOfWords])
-key_adapter: pydantic.TypeAdapter[FileChunkID] = pydantic.TypeAdapter(FileChunkID)
+key_adapter: pydantic.TypeAdapter[MinimalSource] = pydantic.TypeAdapter(
+ MinimalSource
+)
class RAG:
"""A retrieval augmented generation CLI"""
+ def __init__(self):
+ self._bm25_store: BM25[MinimalSource] | None = None
+
def index(
self,
/,
else:
continue
chunks.update(chunk_file(path, max_chunk_size, filetype))
- bm25 = BM25.from_corpus(chunks)
+ bm25 = BM25.from_corpus(chunks, fetch_stopwords())
try:
os.makedirs(index, exist_ok=True)
with open(index + "/index.json", "wb") as f:
logging.getLogger(__name__).error(
f"Failed to write index file {e}"
)
+ exit(1)
+ print(f"Successfully ingested: {data} -> {index}")
- def search(
- self, query: str, /, k: int = 5, index: str = "data/processed"
- ) -> None:
+ def _bm25(self, index: str) -> BM25[MinimalSource]:
+ if self._bm25_store:
+ return self._bm25_store
try:
with open(index + "/index.json") as f:
- bm25: BM25[FileChunkID] = BM25.from_storage(
+ bm25: BM25[MinimalSource] = BM25.from_storage(
(key_adapter.validate_json(k), v)
for k, v in storage_adapter.validate_json(f.read()).items()
)
- for start, length, file in bm25.best_k(query, k):
- print(f"{file} [{start}:{start + length - 1}]")
- except OSError as e:
- print(e)
- pass
- pass
+ self._bm25_store = bm25
+ return bm25
+ except (OSError, ValidationError) as e:
+ logging.getLogger(__name__).error(f"Failed to open index {e}")
+ exit(1)
+
+ def _search(self, query: str, k: int, index) -> list[MinimalSource]:
+ bm25 = self._bm25(index)
+ return bm25.best_k(query, k)
+
+ def search(
+ self, query: str, /, k: int = 5, index: str = "data/processed"
+ ) -> None:
+ bm25 = self._bm25(index)
+ for source in self._search(query, k, index):
+ print(source)
+
+ def search_dataset(
+ self,
+ dataset_path: str,
+ save_directory: str,
+ /,
+ k: int = 5,
+ index: str = "data/processed",
+ ) -> None:
+ bm25 = self._bm25(index)
+ try:
+ with open(dataset_path) as f:
+ dataset = RagDataset.model_validate_json(f.read())
+ except (OSError, ValidationError) as e:
+ logging.getLogger(__name__).error(f"Failed to open dataset {e}")
+ exit(1)
+ result = SearchResults(
+ search_results=[
+ RetrievedQuestion(
+ question_id=question.question_id,
+ question=question.question,
+ retrieved_sources=bm25.best_k(question.question, k),
+ )
+ for question in dataset.rag_questions
+ ],
+ k=k,
+ )
+ try:
+ with open(save_directory +"/"+ os.path.basename(dataset_path), "w") as f:
+ f.write(result.model_dump_json())
+ except (OSError, ValidationError) as e:
+ logging.getLogger(__name__).error(f"Failed to open dataset {e}")
+ exit(1)
import pydantic
+from rag.models import MinimalSource
+
class FileType(Enum):
PYTHON = auto()
)
-type FileChunkID = tuple[int, pydantic.NonNegativeInt, str]
+# type FileChunkID = tuple[int, pydantic.NonNegativeInt, str]
def chunk_by_depth(
def chunk_file(
path: str, max_chunk_size: int, by: FileType
-) -> dict[FileChunkID, str]:
+) -> dict[MinimalSource, str]:
try:
with open(path) as f:
s = f.read()
total_len = 0
res = {}
for chunk in chunks:
- res[(total_len, len(chunk), path)] = chunk
+ res[
+ MinimalSource(
+ first_character_index=total_len,
+ last_character_index=total_len + len(chunk) - 1,
+ file_path=path,
+ )
+ ] = chunk
total_len += len(chunk)
return res
--- /dev/null
+from pydantic import BaseModel, Field
+from typing import List
+import uuid
+
+
+class MinimalSource(BaseModel, frozen=True):
+ file_path: str
+ first_character_index: int
+ last_character_index: int
+
+ def __str__(self) -> str:
+ return "{} [{}:{}]".format(
+ self.file_path,
+ self.first_character_index,
+ self.last_character_index,
+ )
+
+
+class Question(BaseModel):
+ question_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+ question: str
+
+
+class RagDataset(BaseModel):
+ rag_questions: List[Question]
+
+
+class RetrievedQuestion(Question):
+ retrieved_sources: List[MinimalSource]
+
+
+class AnsweredQuestion(RetrievedQuestion):
+ answer: str
+
+
+class SearchResults(BaseModel):
+ search_results: List[RetrievedQuestion]
+ k: int
+
+
+class SearchAnswer(BaseModel):
+ search_results: List[AnsweredQuestion]
+ k: int
yield s[start:].lower()
-def bag_of_words(s: str) -> BagOfWords:
+def bag_of_words(s: str, stopwords: set[str]) -> BagOfWords:
res: BagOfWords = {}
for word in words_normalize(s):
+ if word in stopwords:
+ continue
res[word] = res.get(word, 0) + 1
return res
+def fetch_stopwords(source: str = "stopwords-en.txt") -> set[str]:
+ return {word for word in words_normalize(open(source).read())}
+
+
@dataclass
class BM25[T]:
corpus_words: int
return ((k, v[1]) for k, v in self.corpus.items())
@staticmethod
- def from_corpus(raw_corpus: dict[T, str]) -> "BM25[T]":
+ def from_corpus(
+ raw_corpus: dict[T, str], stopwords: set[str] = set()
+ ) -> "BM25[T]":
return BM25.from_storage(
- (k, bow) for k, s in raw_corpus.items() if (bow := bag_of_words(s))
+ (k, bow)
+ for k, s in raw_corpus.items()
+ if (bow := bag_of_words(s, stopwords))
)
def idf(self, q: str) -> float:
--- /dev/null
+able\r
+about\r
+above\r
+abroad\r
+according\r
+accordingly\r
+across\r
+actually\r
+adj\r
+after\r
+afterwards\r
+again\r
+against\r
+ago\r
+ahead\r
+ain't\r
+all\r
+allow\r
+allows\r
+almost\r
+alone\r
+along\r
+alongside\r
+already\r
+also\r
+although\r
+always\r
+am\r
+amid\r
+amidst\r
+among\r
+amongst\r
+an\r
+and\r
+another\r
+any\r
+anybody\r
+anyhow\r
+anyone\r
+anything\r
+anyway\r
+anyways\r
+anywhere\r
+apart\r
+appear\r
+appreciate\r
+appropriate\r
+are\r
+aren't\r
+around\r
+as\r
+a's\r
+aside\r
+ask\r
+asking\r
+associated\r
+at\r
+available\r
+away\r
+awfully\r
+back\r
+backward\r
+backwards\r
+be\r
+became\r
+because\r
+become\r
+becomes\r
+becoming\r
+been\r
+before\r
+beforehand\r
+begin\r
+behind\r
+being\r
+believe\r
+below\r
+beside\r
+besides\r
+best\r
+better\r
+between\r
+beyond\r
+both\r
+brief\r
+but\r
+by\r
+came\r
+can\r
+cannot\r
+cant\r
+can't\r
+caption\r
+cause\r
+causes\r
+certain\r
+certainly\r
+changes\r
+clearly\r
+c'mon\r
+co\r
+co.\r
+com\r
+come\r
+comes\r
+concerning\r
+consequently\r
+consider\r
+considering\r
+contain\r
+containing\r
+contains\r
+corresponding\r
+could\r
+couldn't\r
+course\r
+c's\r
+currently\r
+dare\r
+daren't\r
+definitely\r
+described\r
+despite\r
+did\r
+didn't\r
+different\r
+directly\r
+do\r
+does\r
+doesn't\r
+doing\r
+done\r
+don't\r
+down\r
+downwards\r
+during\r
+each\r
+edu\r
+eg\r
+eight\r
+eighty\r
+either\r
+else\r
+elsewhere\r
+end\r
+ending\r
+enough\r
+entirely\r
+especially\r
+et\r
+etc\r
+even\r
+ever\r
+evermore\r
+every\r
+everybody\r
+everyone\r
+everything\r
+everywhere\r
+ex\r
+exactly\r
+example\r
+except\r
+fairly\r
+far\r
+farther\r
+few\r
+fewer\r
+fifth\r
+first\r
+five\r
+followed\r
+following\r
+follows\r
+for\r
+forever\r
+former\r
+formerly\r
+forth\r
+forward\r
+found\r
+four\r
+from\r
+further\r
+furthermore\r
+get\r
+gets\r
+getting\r
+given\r
+gives\r
+go\r
+goes\r
+going\r
+gone\r
+got\r
+gotten\r
+greetings\r
+had\r
+hadn't\r
+half\r
+happens\r
+hardly\r
+has\r
+hasn't\r
+have\r
+haven't\r
+having\r
+he\r
+he'd\r
+he'll\r
+hello\r
+help\r
+hence\r
+her\r
+here\r
+hereafter\r
+hereby\r
+herein\r
+here's\r
+hereupon\r
+hers\r
+herself\r
+he's\r
+hi\r
+him\r
+himself\r
+his\r
+hither\r
+hopefully\r
+how\r
+howbeit\r
+however\r
+hundred\r
+i'd\r
+ie\r
+if\r
+ignored\r
+i'll\r
+i'm\r
+immediate\r
+in\r
+inasmuch\r
+inc\r
+inc.\r
+indeed\r
+indicate\r
+indicated\r
+indicates\r
+inner\r
+inside\r
+insofar\r
+instead\r
+into\r
+inward\r
+is\r
+isn't\r
+it\r
+it'd\r
+it'll\r
+its\r
+it's\r
+itself\r
+i've\r
+just\r
+k\r
+keep\r
+keeps\r
+kept\r
+know\r
+known\r
+knows\r
+last\r
+lately\r
+later\r
+latter\r
+latterly\r
+least\r
+less\r
+lest\r
+let\r
+let's\r
+like\r
+liked\r
+likely\r
+likewise\r
+little\r
+look\r
+looking\r
+looks\r
+low\r
+lower\r
+ltd\r
+made\r
+mainly\r
+make\r
+makes\r
+many\r
+may\r
+maybe\r
+mayn't\r
+me\r
+mean\r
+meantime\r
+meanwhile\r
+merely\r
+might\r
+mightn't\r
+mine\r
+minus\r
+miss\r
+more\r
+moreover\r
+most\r
+mostly\r
+mr\r
+mrs\r
+much\r
+must\r
+mustn't\r
+my\r
+myself\r
+name\r
+namely\r
+nd\r
+near\r
+nearly\r
+necessary\r
+need\r
+needn't\r
+needs\r
+neither\r
+never\r
+neverf\r
+neverless\r
+nevertheless\r
+new\r
+next\r
+nine\r
+ninety\r
+no\r
+nobody\r
+non\r
+none\r
+nonetheless\r
+noone\r
+no-one\r
+nor\r
+normally\r
+not\r
+nothing\r
+notwithstanding\r
+novel\r
+now\r
+nowhere\r
+obviously\r
+of\r
+off\r
+often\r
+oh\r
+ok\r
+okay\r
+old\r
+on\r
+once\r
+one\r
+ones\r
+one's\r
+only\r
+onto\r
+opposite\r
+or\r
+other\r
+others\r
+otherwise\r
+ought\r
+oughtn't\r
+our\r
+ours\r
+ourselves\r
+out\r
+outside\r
+over\r
+overall\r
+own\r
+particular\r
+particularly\r
+past\r
+per\r
+perhaps\r
+placed\r
+please\r
+plus\r
+possible\r
+presumably\r
+probably\r
+provided\r
+provides\r
+que\r
+quite\r
+qv\r
+rather\r
+rd\r
+re\r
+really\r
+reasonably\r
+recent\r
+recently\r
+regarding\r
+regardless\r
+regards\r
+relatively\r
+respectively\r
+right\r
+round\r
+said\r
+same\r
+saw\r
+say\r
+saying\r
+says\r
+second\r
+secondly\r
+see\r
+seeing\r
+seem\r
+seemed\r
+seeming\r
+seems\r
+seen\r
+self\r
+selves\r
+sensible\r
+sent\r
+serious\r
+seriously\r
+seven\r
+several\r
+shall\r
+shan't\r
+she\r
+she'd\r
+she'll\r
+she's\r
+should\r
+shouldn't\r
+since\r
+six\r
+so\r
+some\r
+somebody\r
+someday\r
+somehow\r
+someone\r
+something\r
+sometime\r
+sometimes\r
+somewhat\r
+somewhere\r
+soon\r
+sorry\r
+specified\r
+specify\r
+specifying\r
+still\r
+sub\r
+such\r
+sup\r
+sure\r
+take\r
+taken\r
+taking\r
+tell\r
+tends\r
+th\r
+than\r
+thank\r
+thanks\r
+thanx\r
+that\r
+that'll\r
+thats\r
+that's\r
+that've\r
+the\r
+their\r
+theirs\r
+them\r
+themselves\r
+then\r
+thence\r
+there\r
+thereafter\r
+thereby\r
+there'd\r
+therefore\r
+therein\r
+there'll\r
+there're\r
+theres\r
+there's\r
+thereupon\r
+there've\r
+these\r
+they\r
+they'd\r
+they'll\r
+they're\r
+they've\r
+thing\r
+things\r
+think\r
+third\r
+thirty\r
+this\r
+thorough\r
+thoroughly\r
+those\r
+though\r
+three\r
+through\r
+throughout\r
+thru\r
+thus\r
+till\r
+to\r
+together\r
+too\r
+took\r
+toward\r
+towards\r
+tried\r
+tries\r
+truly\r
+try\r
+trying\r
+t's\r
+twice\r
+two\r
+un\r
+under\r
+underneath\r
+undoing\r
+unfortunately\r
+unless\r
+unlike\r
+unlikely\r
+until\r
+unto\r
+up\r
+upon\r
+upwards\r
+us\r
+use\r
+used\r
+useful\r
+uses\r
+using\r
+usually\r
+v\r
+value\r
+various\r
+versus\r
+very\r
+via\r
+viz\r
+vs\r
+want\r
+wants\r
+was\r
+wasn't\r
+way\r
+we\r
+we'd\r
+welcome\r
+well\r
+we'll\r
+went\r
+were\r
+we're\r
+weren't\r
+we've\r
+what\r
+whatever\r
+what'll\r
+what's\r
+what've\r
+when\r
+whence\r
+whenever\r
+where\r
+whereafter\r
+whereas\r
+whereby\r
+wherein\r
+where's\r
+whereupon\r
+wherever\r
+whether\r
+which\r
+whichever\r
+while\r
+whilst\r
+whither\r
+who\r
+who'd\r
+whoever\r
+whole\r
+who'll\r
+whom\r
+whomever\r
+who's\r
+whose\r
+why\r
+will\r
+willing\r
+wish\r
+with\r
+within\r
+without\r
+wonder\r
+won't\r
+would\r
+wouldn't\r
+yes\r
+yet\r
+you\r
+you'd\r
+you'll\r
+your\r
+you're\r
+yours\r
+yourself\r
+yourselves\r
+you've\r
+zero\r
+a\r
+how's\r
+i\r
+when's\r
+why's\r
+b\r
+c\r
+d\r
+e\r
+f\r
+g\r
+h\r
+j\r
+l\r
+m\r
+n\r
+o\r
+p\r
+q\r
+r\r
+s\r
+t\r
+u\r
+uucp\r
+w\r
+x\r
+y\r
+z\r
+I\r
+www\r
+amount\r
+bill\r
+bottom\r
+call\r
+computer\r
+con\r
+couldnt\r
+cry\r
+de\r
+describe\r
+detail\r
+due\r
+eleven\r
+empty\r
+fifteen\r
+fifty\r
+fill\r
+find\r
+fire\r
+forty\r
+front\r
+full\r
+give\r
+hasnt\r
+herse\r
+himse\r
+interest\r
+itse”\r
+mill\r
+move\r
+myse”\r
+part\r
+put\r
+show\r
+side\r
+sincere\r
+sixty\r
+system\r
+ten\r
+thick\r
+thin\r
+top\r
+twelve\r
+twenty\r
+abst\r
+accordance\r
+act\r
+added\r
+adopted\r
+affected\r
+affecting\r
+affects\r
+ah\r
+announce\r
+anymore\r
+apparently\r
+approximately\r
+aren\r
+arent\r
+arise\r
+auth\r
+beginning\r
+beginnings\r
+begins\r
+biol\r
+briefly\r
+ca\r
+date\r
+ed\r
+effect\r
+et-al\r
+ff\r
+fix\r
+gave\r
+giving\r
+heres\r
+hes\r
+hid\r
+home\r
+id\r
+im\r
+immediately\r
+importance\r
+important\r
+index\r
+information\r
+invention\r
+itd\r
+keys\r
+kg\r
+km\r
+largely\r
+lets\r
+line\r
+'ll\r
+means\r
+mg\r
+million\r
+ml\r
+mug\r
+na\r
+nay\r
+necessarily\r
+nos\r
+noted\r
+obtain\r
+obtained\r
+omitted\r
+ord\r
+owing\r
+page\r
+pages\r
+poorly\r
+possibly\r
+potentially\r
+pp\r
+predominantly\r
+present\r
+previously\r
+primarily\r
+promptly\r
+proud\r
+quickly\r
+ran\r
+readily\r
+ref\r
+refs\r
+related\r
+research\r
+resulted\r
+resulting\r
+results\r
+run\r
+sec\r
+section\r
+shed\r
+shes\r
+showed\r
+shown\r
+showns\r
+shows\r
+significant\r
+significantly\r
+similar\r
+similarly\r
+slightly\r
+somethan\r
+specifically\r
+state\r
+states\r
+stop\r
+strongly\r
+substantially\r
+successfully\r
+sufficiently\r
+suggest\r
+thered\r
+thereof\r
+therere\r
+thereto\r
+theyd\r
+theyre\r
+thou\r
+thoughh\r
+thousand\r
+throug\r
+til\r
+tip\r
+ts\r
+ups\r
+usefully\r
+usefulness\r
+'ve\r
+vol\r
+vols\r
+wed\r
+whats\r
+wheres\r
+whim\r
+whod\r
+whos\r
+widely\r
+words\r
+world\r
+youd\r
+youre
\ No newline at end of file
revision = 3
requires-python = ">=3.13"
resolution-markers = [
- "python_full_version >= '3.15' and sys_platform != 'darwin'",
- "python_full_version < '3.15' and sys_platform != 'darwin'",
+ "python_full_version >= '3.15' and sys_platform == 'win32'",
+ "python_full_version >= '3.15' and sys_platform == 'emscripten'",
+ "python_full_version >= '3.15' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'",
+ "python_full_version == '3.14.*' and sys_platform == 'win32'",
+ "python_full_version == '3.14.*' and sys_platform == 'emscripten'",
+ "python_full_version == '3.14.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'",
+ "python_full_version < '3.14' and sys_platform == 'win32'",
+ "python_full_version < '3.14' and sys_platform == 'emscripten'",
+ "python_full_version < '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'",
"python_full_version >= '3.15' and sys_platform == 'darwin'",
- "python_full_version < '3.15' and sys_platform == 'darwin'",
+ "python_full_version == '3.14.*' and sys_platform == 'darwin'",
+ "python_full_version < '3.14' and sys_platform == 'darwin'",
]
[[package]]
[[package]]
name = "fsspec"
-version = "2026.7.0"
+version = "2024.3.1"
source = { registry = "https://download.pytorch.org/whl/cpu" }
-sdist = { url = "https://files.pythonhosted.org/packages/00/78/f34251dadb8f3921264a1d9b8946f5e542014ee2614b285261b4e40e6775/fsspec-2026.7.0.tar.gz", hash = "sha256:c803c40f4cf860b49dea58ee3e1c33cb9c790520e233537e1340049f89b82a88" }
+sdist = { url = "https://files.pythonhosted.org/packages/8b/b8/e3ba21f03c00c27adc9a8cd1cab8adfb37b6024757133924a9a4eab63a83/fsspec-2024.3.1.tar.gz", hash = "sha256:f39780e282d7d117ffb42bb96992f8a90795e4d0fb0f661a70ca39fe9c43ded9" }
wheels = [
- { url = "https://files.pythonhosted.org/packages/fd/3c/6a2bf344106328fd04963664a60b9bb6496fc25df8e962fcdc1367285fb9/fsspec-2026.7.0-py3-none-any.whl", hash = "sha256:b57ddbafedfaef7018c1ecab32aa200a9d7ca26b77965f64e48b70061249d279" },
+ { url = "https://files.pythonhosted.org/packages/93/6d/66d48b03460768f523da62a57a7e14e5e95fdf339d79e996ce3cecda2cdb/fsspec-2024.3.1-py3-none-any.whl", hash = "sha256:918d18d41bf73f0e2b261824baeb1b124bcf771767e3a26425cd7dec3332f512" },
]
[[package]]
version = "2.13.0"
source = { registry = "https://download.pytorch.org/whl/cpu" }
resolution-markers = [
- "python_full_version < '3.15' and sys_platform == 'darwin'",
+ "python_full_version == '3.14.*' and sys_platform == 'darwin'",
+ "python_full_version < '3.14' and sys_platform == 'darwin'",
]
dependencies = [
{ name = "filelock" },
version = "2.13.0+cpu"
source = { registry = "https://download.pytorch.org/whl/cpu" }
resolution-markers = [
- "python_full_version >= '3.15' and sys_platform != 'darwin'",
- "python_full_version < '3.15' and sys_platform != 'darwin'",
+ "python_full_version >= '3.15' and sys_platform == 'win32'",
+ "python_full_version >= '3.15' and sys_platform == 'emscripten'",
+ "python_full_version >= '3.15' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'",
+ "python_full_version == '3.14.*' and sys_platform == 'win32'",
+ "python_full_version == '3.14.*' and sys_platform == 'emscripten'",
+ "python_full_version == '3.14.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'",
+ "python_full_version < '3.14' and sys_platform == 'win32'",
+ "python_full_version < '3.14' and sys_platform == 'emscripten'",
+ "python_full_version < '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'",
"python_full_version >= '3.15' and sys_platform == 'darwin'",
]
dependencies = [
projects / axy / ft / rag.git / commitdiff