andreped/ReferenceBot

1from langchain.docstore.document import Document2from langchain.text_splitter import RecursiveCharacterTextSplitter3 4from knowledge_gpt.core.parsing import File5 6 7def chunk_file(file: File, chunk_size: int, chunk_overlap: int = 0, model_name="gpt-3.5-turbo") -> File:8    """Chunks each document in a file into smaller documents9    according to the specified chunk size and overlap10    where the size is determined by the number of tokens for the specified model.11    """12 13    # split each document into chunks14    chunked_docs = []15    for doc in file.docs:16        text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(17            model_name=model_name,18            chunk_size=chunk_size,19            chunk_overlap=chunk_overlap,20        )21 22        chunks = text_splitter.split_text(doc.page_content)23 24        for i, chunk in enumerate(chunks):25            doc = Document(26                page_content=chunk,27                metadata={28                    "page": doc.metadata.get("page", 1),29                    "chunk": i + 1,30                    "source": f"{doc.metadata.get('page', 1)}-{i + 1}",31                },32            )33            chunked_docs.append(doc)34 35    chunked_file = file.copy()36    chunked_file.docs = chunked_docs37    return chunked_file38