From: Axy Date: Wed, 26 Aug 2026 13:10:57 +0000 (+0200) Subject: silly pydantic models X-Git-Url: https://git.uwuaxy.net/sitemap.xml?a=commitdiff_plain;h=a228208fe19d298dcc4a80e36db76eb6a2e13ad6;p=axy%2Fft%2Frag.git silly pydantic models --- diff --git a/src/rag/__init__.py b/src/rag/__init__.py index 0a11c43..c4587ba 100644 --- a/src/rag/__init__.py +++ b/src/rag/__init__.py @@ -4,22 +4,35 @@ import os 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, /, @@ -38,7 +51,7 @@ class RAG: 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: @@ -54,19 +67,64 @@ class RAG: 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) diff --git a/src/rag/chunking.py b/src/rag/chunking.py index 1da15fe..5e0c80f 100644 --- a/src/rag/chunking.py +++ b/src/rag/chunking.py @@ -4,6 +4,8 @@ from itertools import count import pydantic +from rag.models import MinimalSource + class FileType(Enum): PYTHON = auto() @@ -30,7 +32,7 @@ class FileType(Enum): ) -type FileChunkID = tuple[int, pydantic.NonNegativeInt, str] +# type FileChunkID = tuple[int, pydantic.NonNegativeInt, str] def chunk_by_depth( @@ -68,7 +70,7 @@ 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() @@ -81,6 +83,12 @@ def chunk_file( 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 diff --git a/src/rag/models.py b/src/rag/models.py new file mode 100644 index 0000000..71d4321 --- /dev/null +++ b/src/rag/models.py @@ -0,0 +1,43 @@ +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 diff --git a/src/rag/retrieval.py b/src/rag/retrieval.py index e25214f..80ee944 100644 --- a/src/rag/retrieval.py +++ b/src/rag/retrieval.py @@ -19,13 +19,19 @@ def words_normalize(s: str) -> Iterable[str]: 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 @@ -50,9 +56,13 @@ class BM25[T]: 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: diff --git a/stopwords-en.txt b/stopwords-en.txt new file mode 100644 index 0000000..f9d6929 --- /dev/null +++ b/stopwords-en.txt @@ -0,0 +1,851 @@ +able +about +above +abroad +according +accordingly +across +actually +adj +after +afterwards +again +against +ago +ahead +ain't +all +allow +allows +almost +alone +along +alongside +already +also +although +always +am +amid +amidst +among +amongst +an +and +another +any +anybody +anyhow +anyone +anything +anyway +anyways +anywhere +apart +appear +appreciate +appropriate +are +aren't +around +as +a's +aside +ask +asking +associated +at +available +away +awfully +back +backward +backwards +be +became +because +become +becomes +becoming +been +before +beforehand +begin +behind +being +believe +below +beside +besides +best +better +between +beyond +both +brief +but +by +came +can +cannot +cant +can't +caption +cause +causes +certain +certainly +changes +clearly +c'mon +co +co. +com +come +comes +concerning +consequently +consider +considering +contain +containing +contains +corresponding +could +couldn't +course +c's +currently +dare +daren't +definitely +described +despite +did +didn't +different +directly +do +does +doesn't +doing +done +don't +down +downwards +during +each +edu +eg +eight +eighty +either +else +elsewhere +end +ending +enough +entirely +especially +et +etc +even +ever +evermore +every +everybody +everyone +everything +everywhere +ex +exactly +example +except +fairly +far +farther +few +fewer +fifth +first +five +followed +following +follows +for +forever +former +formerly +forth +forward +found +four +from +further +furthermore +get +gets +getting +given +gives +go +goes +going +gone +got +gotten +greetings +had +hadn't +half +happens +hardly +has +hasn't +have +haven't +having +he +he'd +he'll +hello +help +hence +her +here +hereafter +hereby +herein +here's +hereupon +hers +herself +he's +hi +him +himself +his +hither +hopefully +how +howbeit +however +hundred +i'd +ie +if +ignored +i'll +i'm +immediate +in +inasmuch +inc +inc. +indeed +indicate +indicated +indicates +inner +inside +insofar +instead +into +inward +is +isn't +it +it'd +it'll +its +it's +itself +i've +just +k +keep +keeps +kept +know +known +knows +last +lately +later +latter +latterly +least +less +lest +let +let's +like +liked +likely +likewise +little +look +looking +looks +low +lower +ltd +made +mainly +make +makes +many +may +maybe +mayn't +me +mean +meantime +meanwhile +merely +might +mightn't +mine +minus +miss +more +moreover +most +mostly +mr +mrs +much +must +mustn't +my +myself +name +namely +nd +near +nearly +necessary +need +needn't +needs +neither +never +neverf +neverless +nevertheless +new +next +nine +ninety +no +nobody +non +none +nonetheless +noone +no-one +nor +normally +not +nothing +notwithstanding +novel +now +nowhere +obviously +of +off +often +oh +ok +okay +old +on +once +one +ones +one's +only +onto +opposite +or +other +others +otherwise +ought +oughtn't +our +ours +ourselves +out +outside +over +overall +own +particular +particularly +past +per +perhaps +placed +please +plus +possible +presumably +probably +provided +provides +que +quite +qv +rather +rd +re +really +reasonably +recent +recently +regarding +regardless +regards +relatively +respectively +right +round +said +same +saw +say +saying +says +second +secondly +see +seeing +seem +seemed +seeming +seems +seen +self +selves +sensible +sent +serious +seriously +seven +several +shall +shan't +she +she'd +she'll +she's +should +shouldn't +since +six +so +some +somebody +someday +somehow +someone +something +sometime +sometimes +somewhat +somewhere +soon +sorry +specified +specify +specifying +still +sub +such +sup +sure +take +taken +taking +tell +tends +th +than +thank +thanks +thanx +that +that'll +thats +that's +that've +the +their +theirs +them +themselves +then +thence +there +thereafter +thereby +there'd +therefore +therein +there'll +there're +theres +there's +thereupon +there've +these +they +they'd +they'll +they're +they've +thing +things +think +third +thirty +this +thorough +thoroughly +those +though +three +through +throughout +thru +thus +till +to +together +too +took +toward +towards +tried +tries +truly +try +trying +t's +twice +two +un +under +underneath +undoing +unfortunately +unless +unlike +unlikely +until +unto +up +upon +upwards +us +use +used +useful +uses +using +usually +v +value +various +versus +very +via +viz +vs +want +wants +was +wasn't +way +we +we'd +welcome +well +we'll +went +were +we're +weren't +we've +what +whatever +what'll +what's +what've +when +whence +whenever +where +whereafter +whereas +whereby +wherein +where's +whereupon +wherever +whether +which +whichever +while +whilst +whither +who +who'd +whoever +whole +who'll +whom +whomever +who's +whose +why +will +willing +wish +with +within +without +wonder +won't +would +wouldn't +yes +yet +you +you'd +you'll +your +you're +yours +yourself +yourselves +you've +zero +a +how's +i +when's +why's +b +c +d +e +f +g +h +j +l +m +n +o +p +q +r +s +t +u +uucp +w +x +y +z +I +www +amount +bill +bottom +call +computer +con +couldnt +cry +de +describe +detail +due +eleven +empty +fifteen +fifty +fill +find +fire +forty +front +full +give +hasnt +herse +himse +interest +itse” +mill +move +myse” +part +put +show +side +sincere +sixty +system +ten +thick +thin +top +twelve +twenty +abst +accordance +act +added +adopted +affected +affecting +affects +ah +announce +anymore +apparently +approximately +aren +arent +arise +auth +beginning +beginnings +begins +biol +briefly +ca +date +ed +effect +et-al +ff +fix +gave +giving +heres +hes +hid +home +id +im +immediately +importance +important +index +information +invention +itd +keys +kg +km +largely +lets +line +'ll +means +mg +million +ml +mug +na +nay +necessarily +nos +noted +obtain +obtained +omitted +ord +owing +page +pages +poorly +possibly +potentially +pp +predominantly +present +previously +primarily +promptly +proud +quickly +ran +readily +ref +refs +related +research +resulted +resulting +results +run +sec +section +shed +shes +showed +shown +showns +shows +significant +significantly +similar +similarly +slightly +somethan +specifically +state +states +stop +strongly +substantially +successfully +sufficiently +suggest +thered +thereof +therere +thereto +theyd +theyre +thou +thoughh +thousand +throug +til +tip +ts +ups +usefully +usefulness +'ve +vol +vols +wed +whats +wheres +whim +whod +whos +widely +words +world +youd +youre \ No newline at end of file diff --git a/uv.lock b/uv.lock index f82a919..9b79cb5 100644 --- a/uv.lock +++ b/uv.lock @@ -2,10 +2,18 @@ version = 1 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]] @@ -167,11 +175,11 @@ wheels = [ [[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]] @@ -940,7 +948,8 @@ name = "torch" 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" }, @@ -962,8 +971,15 @@ name = "torch" 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 = [