Add chunker.py
This commit is contained in:
+15
-27
@@ -1,46 +1,34 @@
|
||||
"""
|
||||
Chunking utilities for the RAG agent.
|
||||
|
||||
Uses RecursiveCharacterTextSplitter from LangChain with chunk_size=500 and overlap=100.
|
||||
Uses RecursiveCharacterTextSplitter from LangChain with chunk_size=500 and chunk_overlap=100.
|
||||
"""
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
|
||||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||
|
||||
# Default splitter configuration
|
||||
CHUNK_SIZE = 500
|
||||
OVERLAP = 100
|
||||
DEFAULT_SPLITTER = RecursiveCharacterTextSplitter(
|
||||
chunk_size=500,
|
||||
chunk_overlap=100,
|
||||
)
|
||||
|
||||
splitter = RecursiveCharacterTextSplitter(chunk_size=CHUNK_SIZE, chunk_overlap=OVERLAP)
|
||||
|
||||
def split_text(text: str) -> List[str]:
|
||||
def split_text(text: str):
|
||||
"""Split a large string into chunks.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
text: str
|
||||
The raw document content.
|
||||
Text to split.
|
||||
|
||||
Returns
|
||||
-------
|
||||
List[str]
|
||||
A list of chunk strings.
|
||||
list[str]
|
||||
List of chunk strings.
|
||||
"""
|
||||
return splitter.split_text(text)
|
||||
return DEFAULT_SPLITTER.split_text(text)
|
||||
|
||||
# Helper to read all files from a directory and split them
|
||||
|
||||
def load_and_split(directory: Path) -> List[tuple]:
|
||||
"""Load text files from *directory* and split into chunks.
|
||||
|
||||
Returns a list of tuples (chunk, metadata).
|
||||
Metadata contains the source file path.
|
||||
"""
|
||||
chunks = []
|
||||
for file_path in directory.rglob("*.txt"):
|
||||
content = file_path.read_text(encoding="utf-8")
|
||||
for chunk in split_text(content):
|
||||
chunks.append((chunk, {"source": str(file_path)}))
|
||||
return chunks
|
||||
# Example usage (not executed in tests)
|
||||
if __name__ == "__main__":
|
||||
sample = "\n".join([f"Line {i}" for i in range(1000)])
|
||||
chunks = split_text(sample)
|
||||
print(f"Generated {len(chunks)} chunks")
|
||||
Reference in New Issue
Block a user