feat: add chunking methods and ingestion process for Elasticsearch
This commit is contained in:
pseco
parent
f2482cae19
commit
26603a9f45
|
|
@ -0,0 +1,396 @@
|
|||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 15,
|
||||
"id": "0a8abbfa",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"True"
|
||||
]
|
||||
},
|
||||
"execution_count": 15,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"import os\n",
|
||||
"import re\n",
|
||||
"import uuid\n",
|
||||
"from dataclasses import dataclass\n",
|
||||
"from typing import Iterable, List, Dict, Any, Callable, Protocol\n",
|
||||
"\n",
|
||||
"import torch\n",
|
||||
"import torch.nn.functional as F\n",
|
||||
"from loguru import logger\n",
|
||||
"from transformers import AutoTokenizer, AutoModel\n",
|
||||
"from elasticsearch import Elasticsearch\n",
|
||||
"from elasticsearch.helpers import bulk\n",
|
||||
"import nltk\n",
|
||||
"from nltk.tokenize import sent_tokenize\n",
|
||||
"nltk.download(\"punkt\", quiet=True)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "77f6c552",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Domain model"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"id": "c4cd2bc2",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"@dataclass(frozen=True)\n",
|
||||
"class Chunk:\n",
|
||||
" doc_id: str\n",
|
||||
" chunk_id: int\n",
|
||||
" text: str\n",
|
||||
" source: str\n",
|
||||
" metadata: Dict[str, Any]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "5cd700bd",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Utilities"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"id": "84e834d9",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def clean_text(text: str) -> str:\n",
|
||||
" text = text.replace(\"\\u00a0\", \" \")\n",
|
||||
" text = re.sub(r\"\\s+\", \" \", text).strip()\n",
|
||||
" return text"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"id": "4ebdc5f5",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"class ChunkingStrategy(Protocol):\n",
|
||||
" def __call__(self, text: str, **kwargs) -> List[str]:\n",
|
||||
" ..."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "82209fc0",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Chunking strategies"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 23,
|
||||
"id": "9f360449",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def fixed_size_token_chunking(\n",
|
||||
" text: str,\n",
|
||||
" embedding_model_name: str = os.getenv(\"EMBEDDING_MODEL_NAME\"),\n",
|
||||
" chunk_size: int = 1200,\n",
|
||||
" overlap: int = 200,\n",
|
||||
") -> List[str]:\n",
|
||||
"\n",
|
||||
" if chunk_size <= overlap:\n",
|
||||
" raise ValueError(\"chunk_size must be greater than overlap\")\n",
|
||||
"\n",
|
||||
" tokenizer = AutoTokenizer.from_pretrained(embedding_model_name, use_fast=True)\n",
|
||||
" token_ids = tokenizer.encode(text, add_special_tokens=False)\n",
|
||||
"\n",
|
||||
" chunks: List[str] = []\n",
|
||||
" start = 0\n",
|
||||
" n = len(token_ids)\n",
|
||||
"\n",
|
||||
" while start < n:\n",
|
||||
" end = min(start + chunk_size, n)\n",
|
||||
" chunk_ids = token_ids[start:end]\n",
|
||||
" chunks.append(tokenizer.decode(chunk_ids, skip_special_tokens=True))\n",
|
||||
"\n",
|
||||
" if end == n:\n",
|
||||
" break\n",
|
||||
"\n",
|
||||
" start = end - overlap\n",
|
||||
"\n",
|
||||
" return chunks\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def semantic_chunking(\n",
|
||||
" text: str,\n",
|
||||
" embedding_model_name: str = os.getenv(\"EMBEDDING_MODEL_NAME\"),\n",
|
||||
" similarity_threshold: float = 0.78,\n",
|
||||
" max_sentences_per_chunk: int = 12,\n",
|
||||
") -> List[str]:\n",
|
||||
" sentences = [s.strip() for s in sent_tokenize(text) if s.strip()]\n",
|
||||
" if not sentences:\n",
|
||||
" return []\n",
|
||||
" print(f\"Semantic chunking: {len(sentences)} sentences found\")\n",
|
||||
" device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
|
||||
"\n",
|
||||
" tokenizer = AutoTokenizer.from_pretrained(embedding_model_name)\n",
|
||||
" model = AutoModel.from_pretrained(embedding_model_name).to(device)\n",
|
||||
" model.eval()\n",
|
||||
"\n",
|
||||
" with torch.no_grad():\n",
|
||||
" enc = tokenizer(sentences, padding=True, truncation=True, return_tensors=\"pt\").to(device)\n",
|
||||
" out = model(**enc)\n",
|
||||
" mask = enc[\"attention_mask\"].unsqueeze(-1)\n",
|
||||
" vecs = (out.last_hidden_state * mask).sum(1) / mask.sum(1).clamp(min=1e-9)\n",
|
||||
" vecs = F.normalize(vecs, p=2, dim=1)\n",
|
||||
"\n",
|
||||
" chunks: List[List[str]] = [[sentences[0]]]\n",
|
||||
"\n",
|
||||
" for i in range(1, len(sentences)):\n",
|
||||
" sim = float((vecs[i - 1] * vecs[i]).sum())\n",
|
||||
" if sim < similarity_threshold or len(chunks[-1]) >= max_sentences_per_chunk:\n",
|
||||
" chunks.append([])\n",
|
||||
" chunks[-1].append(sentences[i])\n",
|
||||
"\n",
|
||||
" return [\" \".join(chunk) for chunk in chunks if chunk]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"id": "bc7267d7",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"CHUNKING_REGISTRY: Dict[str, ChunkingStrategy] = {\n",
|
||||
" \"fixed\": fixed_size_token_chunking,\n",
|
||||
" \"semantic\": semantic_chunking,\n",
|
||||
"}"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 7,
|
||||
"id": "87f2f70c",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def build_chunks(\n",
|
||||
" doc_text: str,\n",
|
||||
" source: str,\n",
|
||||
" metadata: Dict[str, Any],\n",
|
||||
" chunking_strategy: str = \"fixed\",\n",
|
||||
" **chunking_kwargs,\n",
|
||||
") -> List[Chunk]:\n",
|
||||
"\n",
|
||||
" if chunking_strategy not in CHUNKING_REGISTRY:\n",
|
||||
" raise ValueError(\n",
|
||||
" f\"Unknown chunking strategy '{chunking_strategy}'. \"\n",
|
||||
" f\"Available: {list(CHUNKING_REGISTRY.keys())}\"\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" doc_id = metadata.get(\"doc_id\") or str(uuid.uuid4())\n",
|
||||
" cleaned = clean_text(doc_text)\n",
|
||||
"\n",
|
||||
" chunking_fn = CHUNKING_REGISTRY[chunking_strategy]\n",
|
||||
"\n",
|
||||
" parts = chunking_fn(cleaned, **chunking_kwargs)\n",
|
||||
"\n",
|
||||
" return [\n",
|
||||
" Chunk(\n",
|
||||
" doc_id=doc_id,\n",
|
||||
" chunk_id=i,\n",
|
||||
" text=part,\n",
|
||||
" source=source,\n",
|
||||
" metadata={**metadata, \"doc_id\": doc_id},\n",
|
||||
" )\n",
|
||||
" for i, part in enumerate(parts)\n",
|
||||
" if part.strip()\n",
|
||||
" ]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "ba5649e9",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Ingestion in elasticsearch"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "ff03c689",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def index_chunks(\n",
|
||||
" es: Elasticsearch,\n",
|
||||
" index_name: str,\n",
|
||||
" model: SentenceTransformer,\n",
|
||||
" chunks: List[Chunk],\n",
|
||||
" batch_size: int = 64,\n",
|
||||
") -> None:\n",
|
||||
" def actions() -> Iterable[Dict[str, Any]]:\n",
|
||||
" # Embed in batches for speed\n",
|
||||
" for i in range(0, len(chunks), batch_size):\n",
|
||||
" batch = chunks[i:i + batch_size]\n",
|
||||
" texts = [c.text for c in batch]\n",
|
||||
" vectors = model.encode(texts, normalize_embeddings=True).tolist()\n",
|
||||
"\n",
|
||||
" for c, v in zip(batch, vectors):\n",
|
||||
" yield {\n",
|
||||
" \"_op_type\": \"index\",\n",
|
||||
" \"_index\": index_name,\n",
|
||||
" \"_id\": f\"{c.doc_id}:{c.chunk_id}\",\n",
|
||||
" \"_source\": {\n",
|
||||
" \"doc_id\": c.doc_id,\n",
|
||||
" \"chunk_id\": c.chunk_id,\n",
|
||||
" \"text\": c.text,\n",
|
||||
" \"source\": c.source,\n",
|
||||
" \"metadata\": c.metadata,\n",
|
||||
" \"embedding\": v,\n",
|
||||
" },\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" bulk(es, actions(), request_timeout=120)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "7bcf0c87",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"es_url = os.environ.get(\"ES_URL\", \"http://localhost:9200\")\n",
|
||||
"es_user = os.environ.get(\"ES_USER\")\n",
|
||||
"es_pass = os.environ.get(\"ES_PASS\")\n",
|
||||
"index_name = \"my_docs_v1\"\n",
|
||||
"\n",
|
||||
"es = Elasticsearch(\n",
|
||||
" es_url,\n",
|
||||
" basic_auth=(es_user, es_pass) if es_user and es_pass else None,\n",
|
||||
" request_timeout=60,\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"# Pick a model with dims matching your index mapping.\n",
|
||||
"model = SentenceTransformer(\"sentence-transformers/all-mpnet-base-v2\") # 768 dims\n",
|
||||
"\n",
|
||||
"# Example document\n",
|
||||
"doc_text = \"\"\"\n",
|
||||
"This is a sample document. Replace this with your PDF/HTML extraction output.\n",
|
||||
"\"\"\"\n",
|
||||
"chunks = build_chunks(\n",
|
||||
" doc_text=doc_text,\n",
|
||||
" source=\"local_demo\",\n",
|
||||
" metadata={\"title\": \"Demo\", \"doc_id\": \"demo-001\"},\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"index_chunks(es, index_name, model, chunks)\n",
|
||||
"es.indices.refresh(index=index_name)\n",
|
||||
"print(f\"Indexed {len(chunks)} chunks into {index_name}.\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 24,
|
||||
"id": "b1ba8e85",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"application/vnd.jupyter.widget-view+json": {
|
||||
"model_id": "b763a493689549a180ab815567520c0a",
|
||||
"version_major": 2,
|
||||
"version_minor": 0
|
||||
},
|
||||
"text/plain": [
|
||||
"Loading weights: 0%| | 0/310 [00:00<?, ?it/s]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# Example document\n",
|
||||
"doc_text = \"\"\"\n",
|
||||
"This is a sample document. Replace this with your PDF/HTML extraction output.\n",
|
||||
"\"\"\"\n",
|
||||
"chunks = build_chunks(\n",
|
||||
" doc_text=doc_text,\n",
|
||||
" source=\"local_demo\",\n",
|
||||
" metadata={\"title\": \"Demo\", \"doc_id\": \"demo-001\"},\n",
|
||||
" chunking_strategy=\"semantic\"\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 25,
|
||||
"id": "b2c52b38",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[Chunk(doc_id='demo-001', chunk_id=0, text='This is a sample document. Replace this with your PDF/HTML extraction output.', source='local_demo', metadata={'title': 'Demo', 'doc_id': 'demo-001'})]"
|
||||
]
|
||||
},
|
||||
"execution_count": 25,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"chunks"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "daa57061",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": []
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "assistance-engine",
|
||||
"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.11.13"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
Loading…
Reference in New Issue