{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Chunk by section\n",
    "import re\n",
    "\n",
    "\n",
    "def chunk_by_section(document_text):\n",
    "    pattern = r\"\\n## \"\n",
    "    return re.split(pattern, document_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# BM25Index implementation\n",
    "import math\n",
    "from collections import Counter\n",
    "from typing import Callable, Optional, Any, List, Dict, Tuple\n",
    "\n",
    "\n",
    "class BM25Index:\n",
    "    def __init__(\n",
    "        self,\n",
    "        k1: float = 1.5,\n",
    "        b: float = 0.75,\n",
    "        tokenizer: Optional[Callable[[str], List[str]]] = None,\n",
    "    ):\n",
    "        self.documents: List[Dict[str, Any]] = []\n",
    "        self._corpus_tokens: List[List[str]] = []\n",
    "        self._doc_len: List[int] = []\n",
    "        self._doc_freqs: Dict[str, int] = {}\n",
    "        self._avg_doc_len: float = 0.0\n",
    "        self._idf: Dict[str, float] = {}\n",
    "        self._index_built: bool = False\n",
    "\n",
    "        self.k1 = k1\n",
    "        self.b = b\n",
    "        self._tokenizer = tokenizer if tokenizer else self._default_tokenizer\n",
    "\n",
    "    def _default_tokenizer(self, text: str) -> List[str]:\n",
    "        text = text.lower()\n",
    "        tokens = re.split(r\"\\W+\", text)\n",
    "        return [token for token in tokens if token]\n",
    "\n",
    "    def _update_stats_add(self, doc_tokens: List[str]):\n",
    "        self._doc_len.append(len(doc_tokens))\n",
    "\n",
    "        seen_in_doc = set()\n",
    "        for token in doc_tokens:\n",
    "            if token not in seen_in_doc:\n",
    "                self._doc_freqs[token] = self._doc_freqs.get(token, 0) + 1\n",
    "                seen_in_doc.add(token)\n",
    "\n",
    "        self._index_built = False\n",
    "\n",
    "    def _calculate_idf(self):\n",
    "        N = len(self.documents)\n",
    "        self._idf = {}\n",
    "        for term, freq in self._doc_freqs.items():\n",
    "            idf_score = math.log(((N - freq + 0.5) / (freq + 0.5)) + 1)\n",
    "            self._idf[term] = idf_score\n",
    "\n",
    "    def _build_index(self):\n",
    "        if not self.documents:\n",
    "            self._avg_doc_len = 0.0\n",
    "            self._idf = {}\n",
    "            self._index_built = True\n",
    "            return\n",
    "\n",
    "        self._avg_doc_len = sum(self._doc_len) / len(self.documents)\n",
    "        self._calculate_idf()\n",
    "        self._index_built = True\n",
    "\n",
    "    def add_document(self, document: Dict[str, Any]):\n",
    "        if not isinstance(document, dict):\n",
    "            raise TypeError(\"Document must be a dictionary.\")\n",
    "        if \"content\" not in document:\n",
    "            raise ValueError(\n",
    "                \"Document dictionary must contain a 'content' key.\"\n",
    "            )\n",
    "\n",
    "        content = document.get(\"content\", \"\")\n",
    "        if not isinstance(content, str):\n",
    "            raise TypeError(\"Document 'content' must be a string.\")\n",
    "\n",
    "        doc_tokens = self._tokenizer(content)\n",
    "\n",
    "        self.documents.append(document)\n",
    "        self._corpus_tokens.append(doc_tokens)\n",
    "        self._update_stats_add(doc_tokens)\n",
    "\n",
    "    def _compute_bm25_score(\n",
    "        self, query_tokens: List[str], doc_index: int\n",
    "    ) -> float:\n",
    "        score = 0.0\n",
    "        doc_term_counts = Counter(self._corpus_tokens[doc_index])\n",
    "        doc_length = self._doc_len[doc_index]\n",
    "\n",
    "        for token in query_tokens:\n",
    "            if token not in self._idf:\n",
    "                continue\n",
    "\n",
    "            idf = self._idf[token]\n",
    "            term_freq = doc_term_counts.get(token, 0)\n",
    "\n",
    "            numerator = idf * term_freq * (self.k1 + 1)\n",
    "            denominator = term_freq + self.k1 * (\n",
    "                1 - self.b + self.b * (doc_length / self._avg_doc_len)\n",
    "            )\n",
    "            score += numerator / (denominator + 1e-9)\n",
    "\n",
    "        return score\n",
    "\n",
    "    def search(\n",
    "        self,\n",
    "        query_text: str,\n",
    "        k: int = 1,\n",
    "        score_normalization_factor: float = 0.1,\n",
    "    ) -> List[Tuple[Dict[str, Any], float]]:\n",
    "        if not self.documents:\n",
    "            return []\n",
    "\n",
    "        if not isinstance(query_text, str):\n",
    "            raise TypeError(\"Query text must be a string.\")\n",
    "\n",
    "        if k <= 0:\n",
    "            raise ValueError(\"k must be a positive integer.\")\n",
    "\n",
    "        if not self._index_built:\n",
    "            self._build_index()\n",
    "\n",
    "        if self._avg_doc_len == 0:\n",
    "            return []\n",
    "\n",
    "        query_tokens = self._tokenizer(query_text)\n",
    "        if not query_tokens:\n",
    "            return []\n",
    "\n",
    "        raw_scores = []\n",
    "        for i in range(len(self.documents)):\n",
    "            raw_score = self._compute_bm25_score(query_tokens, i)\n",
    "            if raw_score > 1e-9:\n",
    "                raw_scores.append((raw_score, self.documents[i]))\n",
    "\n",
    "        raw_scores.sort(key=lambda item: item[0], reverse=True)\n",
    "\n",
    "        normalized_results = []\n",
    "        for raw_score, doc in raw_scores[:k]:\n",
    "            normalized_score = math.exp(-score_normalization_factor * raw_score)\n",
    "            normalized_results.append((doc, normalized_score))\n",
    "\n",
    "        normalized_results.sort(key=lambda item: item[1])\n",
    "\n",
    "        return normalized_results\n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        return len(self.documents)\n",
    "\n",
    "    def __repr__(self) -> str:\n",
    "        return f\"BM25VectorStore(count={len(self)}, k1={self.k1}, b={self.b}, index_built={self._index_built})\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"./report.md\", \"r\") as f:\n",
    "    text = f.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. Chunk the text by section"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2. Create a BM25 store and add each chunk to it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 3. Search the store"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}