File size: 1,049 Bytes
db774b8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 | from pathlib import Path
from llama_index.core import SimpleDirectoryReader
from llama_index.core.node_parser import SentenceSplitter
from src.models import Chunk
from src.config import CHUNK_SIZE, CHUNK_OVERLAP
_splitter = SentenceSplitter(chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP)
def load_documents(docs_dir: str | Path) -> list[Chunk]:
docs_dir = Path(docs_dir)
if not docs_dir.exists():
raise FileNotFoundError(f"Docs directory not found: {docs_dir}")
reader = SimpleDirectoryReader(
input_dir=str(docs_dir),
required_exts=[".txt", ".md", ".pdf"],
filename_as_id=True,
)
documents = reader.load_data()
nodes = _splitter.get_nodes_from_documents(documents)
chunks = []
for i, node in enumerate(nodes):
source = node.metadata.get("file_name", "unknown")
chunk = Chunk(
id=f"chunk_{i}",
text=node.text,
source=source,
metadata=node.metadata,
)
chunks.append(chunk)
return chunks
|