From: Axy Date: Thu, 27 Aug 2026 20:59:45 +0000 (+0200) Subject: Only missing flake8 and readme X-Git-Url: https://git.uwuaxy.net/sitemap.xml?a=commitdiff_plain;h=d61026dd6cbbd4d6bd73a6833a5f7f70186ca0be;p=axy%2Fft%2Frag.git Only missing flake8 and readme --- diff --git a/src/rag/__init__.py b/src/rag/__init__.py index f0b6e63..681a1c3 100644 --- a/src/rag/__init__.py +++ b/src/rag/__init__.py @@ -1,16 +1,16 @@ import logging import os from collections.abc import Callable -from sys import stderr import fire import pydantic -from pydantic_core import ValidationError from tqdm import tqdm from rag.answering import InferenceBackend from rag.chunking import FileType, chunk_file +from rag.evaluate import sources_overlap from rag.models import ( + AnswerSet, RagDataset, RetrievedQuestion, SearchAnswers, @@ -41,7 +41,9 @@ class RAG: with open(path) as f: return cb(f.read()) except Exception as e: - logging.getLogger(__name__).error(f"Failed to read {name}: {e}") + logging.getLogger(__name__).error( + f"Failed to read {name} ({path}): {e}" + ) exit(1) @staticmethod @@ -55,7 +57,9 @@ class RAG: with open(path, "w") as f: f.write(data) except Exception as e: - logging.getLogger(__name__).error(f"Failed to write {name}: {e}") + logging.getLogger(__name__).error( + f"Failed to write {name} ({path}): {e}" + ) exit(1) def index( @@ -65,9 +69,14 @@ class RAG: index: str = "data/processed", max_chunk_size: int = 2000, ) -> None: + if max_chunk_size < 1: + logging.getLogger(__name__).error( + f"Invalid max chunk size {max_chunk_size}" + ) + exit(1) chunks = {} for dirpath, _, files in tqdm( - list(os.walk(data)), desc="Collecting files" + list(os.walk(data)), desc="Chunking files" ): for file in files: path = dirpath + "/" + file @@ -114,6 +123,9 @@ class RAG: def search( self, query: str, /, k: int = 5, index: str = "data/processed" ) -> None: + if k < 1: + logging.getLogger(__name__).error(f"Invalid k {k}") + exit(1) for source in self._search(query, k, index): print(source) @@ -125,6 +137,9 @@ class RAG: k: int = 5, index: str = "data/processed", ) -> None: + 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) result = SearchResults( @@ -159,6 +174,14 @@ class RAG: index: str = "data/processed", max_tokens: int = 100, ) -> None: + if k < 1: + logging.getLogger(__name__).error(f"Invalid k {k}") + exit(1) + if max_tokens < 1: + logging.getLogger(__name__).error( + f"Invalid max tokens {max_tokens}" + ) + exit(1) sources = self._search(query, k, index) inference = self._inference() print("Sources:") @@ -176,10 +199,13 @@ class RAG: student_search_result_path: str, save_directory: str, /, - index: str = "data/processed", max_tokens: int = 100, ) -> None: - bm25 = self._bm25(index) + if max_tokens < 1: + logging.getLogger(__name__).error( + f"Invalid max tokens {max_tokens}" + ) + exit(1) dataset = self._readf( student_search_result_path, SearchResults.model_validate_json, @@ -201,3 +227,59 @@ class RAG: + os.path.basename(student_search_result_path), result.model_dump_json(), ) + + def _evaluate( + self, answers: SearchResults, references: AnswerSet, k: int, iou: float + ) -> None: + count = len(references.rag_questions) + success = sum( + 1 + for answer, reference in zip( + answers.search_results, + references.rag_questions, + strict=False, + ) + if any( + sources_overlap(a, b, iou) + for i, a in enumerate(answer.retrieved_sources) + for b in reference.sources + if i < k + ) + ) + print( + f"recall@{k}: {success / count * 100.0:.2f}% ({success} / {count})" + ) + + def evaluate( + self, + student_search_result_path: str, + dataset_path: str, + /, + k: int = 5, + iou: float = 0.05, + ) -> None: + if k < 1: + logging.getLogger(__name__).error(f"Invalid k {k}") + exit(1) + answers = self._readf( + student_search_result_path, + SearchResults.model_validate_json, + name="student search results", + ) + references = self._readf( + dataset_path, + AnswerSet.model_validate_json, + name="reference dataset", + ) + if len(answers.search_results) != len(references.rag_questions): + print("Mismatch between search result and reference set lenghts!") + exit(1) + if len(answers.search_results) == 0: + print("Empty dataset, grade would not make sense.") + exit(0) + prev_i = 1 + i = 1 + while i < k: + self._evaluate(answers, references, i, iou) + prev_i, i = i, i + prev_i + self._evaluate(answers, references, k, iou) diff --git a/src/rag/answering.py b/src/rag/answering.py index f201793..008d13d 100644 --- a/src/rag/answering.py +++ b/src/rag/answering.py @@ -2,9 +2,6 @@ import logging from collections.abc import Callable from typing import Any, cast -from transformers import AutoModelForCausalLM, AutoTokenizer -from transformers.generation import BaseStreamer # type:ignore - from rag.models import AnsweredQuestion, RetrievedQuestion, Source @@ -17,28 +14,9 @@ class InferenceBackend: do not mention them. End your answer with a newline.""" - class _Streamer(BaseStreamer): - def __init__( - self, backend: "InferenceBackend", cb: Callable[[str], None] - ) -> None: - self.backend = backend - self.cb = cb - self.skip = True - - def put(self, value: Any) -> None: - if self.skip: - self.skip = False - return - s: str = cast( - str, - self.backend.tokenizer.decode(value, skip_special_tokens=True), - ) - self.cb(s.rstrip("\n")) - - def end(self) -> None: - self.cb("\n") - def __init__(self, model: str) -> None: + from transformers import AutoModelForCausalLM, AutoTokenizer + self.model: Any = AutoModelForCausalLM.from_pretrained( model, device_map="auto" ) @@ -86,6 +64,31 @@ class InferenceBackend: max_tokens: int, cb: Callable[[str], None], ) -> str: + from transformers.generation import BaseStreamer # type: ignore + + class Streamer(BaseStreamer): + def __init__( + self, backend: "InferenceBackend", cb: Callable[[str], None] + ) -> None: + self.backend = backend + self.cb = cb + self.skip = True + + def put(self, value: Any) -> None: + if self.skip: + self.skip = False + return + s: str = cast( + str, + self.backend.tokenizer.decode( + value, skip_special_tokens=True + ), + ) + self.cb(s.rstrip("\n")) + + def end(self) -> None: + self.cb("\n") + prompt = self._prompt(question) model_inputs = self.tokenizer([prompt], return_tensors="pt").to( self.model.device @@ -95,7 +98,7 @@ class InferenceBackend: max_new_tokens=max_tokens, stop_strings="\n", tokenizer=self.tokenizer, - streamer=self._Streamer(self, cb), + streamer=Streamer(self, cb), ) full_response = self.tokenizer.batch_decode( generated_ids, diff --git a/src/rag/chunking.py b/src/rag/chunking.py index aabaf73..e14c6c3 100644 --- a/src/rag/chunking.py +++ b/src/rag/chunking.py @@ -43,7 +43,10 @@ def chunk_by_depth( chunks.pop() chunks.append(line) if s.count("\n") <= 1: - chunks = s.split() + chunks = s.split(" ") + chunks = [s for s in chunks if s] + for i in range(len(chunks) - 1): + chunks[i] = chunks[i] + " " if len(chunks) <= 1: chunks = [ s[i : min(len(chunks), i + max_chunk_size)] diff --git a/src/rag/evaluate.py b/src/rag/evaluate.py new file mode 100644 index 0000000..090ff5a --- /dev/null +++ b/src/rag/evaluate.py @@ -0,0 +1,14 @@ +from rag.models import Source + + +def sources_overlap(a: Source, b: Source, min_iou: float) -> bool: + 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 + if union < 1: + union = 1 + intersection_start = max(a.first_character_index, b.first_character_index) + intersection_end = min(a.last_character_index, b.last_character_index) + intersection = intersection_end - intersection_start + iou = intersection / union + return iou >= min_iou diff --git a/src/rag/models.py b/src/rag/models.py index cb67b17..4ba9f78 100644 --- a/src/rag/models.py +++ b/src/rag/models.py @@ -41,3 +41,16 @@ class SearchResults(BaseModel): class SearchAnswers(BaseModel): search_results: list[AnsweredQuestion] k: int + + +class DatasetAnswer(BaseModel): + question_id: str + question: str + answer: str + sources: list[Source] + difficulty: str + is_valid: bool + + +class AnswerSet(BaseModel): + rag_questions: list[DatasetAnswer]