【LangChain系列】第二节：文档拆分

import os
from langchain_openai import OpenAI
from dotenv import load_dotenv, find_dotenv

_ = load_dotenv(find_dotenv())

client = OpenAI(
    api_key=os.getenv("OPENAI_API_KEY")
)

from langchain_text_splitters import (
    CharacterTextSplitter,
    RecursiveCharacterTextSplitter,
)

chunk_size = 26
chunk_overlap = 4

r_splitter = RecursiveCharacterTextSplitter(
    chunk_size=chunk_size, chunk_overlap=chunk_overlap
)
c_splitter = CharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)

text1 = "abcdefghijklmnopqrstuvwxyz"
print(r_splitter.split_text(text1))  
# Output: ['abcdefghijklmnopqrstuvwxyz']

text2 = "abcdefghijklmnopqrstuvwxyzabcdefg"
print(r_splitter.split_text(text2))  
# Output: ['abcdefghijklmnopqrstuvwxyz', 'wxyzabcdefg']

text3 = "a b c d e f g h i j k l m n o p q r s t u v w x y z"
print(r_splitter.split_text(text3))  
# Output: ['a b c d e f g h i j k l m', 'l m n o p q r s t u v w x', 'w x y z']
print(c_splitter.split_text(text3))  
# Output: ['a b c d e f g h i j k l m n o p q r s t u v w x y z']

# Set the separator for CharacterTextSplitter
c_splitter = CharacterTextSplitter(
    chunk_size=chunk_size, chunk_overlap=chunk_overlap, separator=" "
)
print(c_splitter.split_text(text3))  
# Output: ['a b c d e f g h i j k l m', 'l m n o p q r s t u v w x', 'w x y z']

some_text = """When writing documents, writers will use document structure to group content. \
This can convey to the reader, which idea's are related. For example, closely related ideas \
are in sentances. Similar ideas are in paragraphs. Paragraphs form a document. \\n\\n  \
Paragraphs are often delimited with a carriage return or two carriage returns. \
Carriage returns are the "backslash n" you see embedded in this string. \
Sentences have a period at the end, but also, have a space.\
and words are separated by space."""

c_splitter = CharacterTextSplitter(chunk_size=450, chunk_overlap=0, separator=" ")
r_splitter = RecursiveCharacterTextSplitter(
    chunk_size=450, chunk_overlap=0, separators=["\n\n", "\n", " ", ""]
)

chunks = c_splitter.split_text(some_text)
print("Chunks: ", chunks)
print("Length of chunks: ", len(chunks))
# Chunks:  ['When writing documents, writers will use document structure to group content. This can convey to the reader, which idea\'s are related. For example, closely related ideas are in sentances. Similar ideas are in paragraphs. Paragraphs form a document. \n\n Paragraphs are often delimited with a carriage return or two carriage returns. Carriage returns are the "backslash n" you see embedded in this string. Sentences have a period at the end, but also,', 'have a space.and words are separated by space.']
# Length of chunks:  2

chunks = r_splitter.split_text(some_text)
print("Chunks: ", chunks)
print("Length of chunks: ", len(chunks))
# Chunks:  ["When writing documents, writers will use document structure to group content. This can convey to the reader, which idea's are related. For example, closely related ideas are in sentances. Similar ideas are in paragraphs. Paragraphs form a document.", 'Paragraphs are often delimited with a carriage return or two carriage returns. Carriage returns are the "backslash n" you see embedded in this string. Sentences have a period at the end, but also, have a space.and words are separated by space.']
# Length of chunks:  2

from langchain.document_loaders import PyPDFLoader, NotionDirectoryLoader

# Load a PDF document
loader = PyPDFLoader("docs/cs229_lectures/MachineLearning-Lecture01.pdf")
pages = loader.load()

text_splitter = CharacterTextSplitter(
    separator="\n", chunk_size=1000, chunk_overlap=150, length_function=len
)
docs = text_splitter.split_documents(pages)

print("Pages in the original document: ", len(pages))
print("Length of chunks after splitting pages: ", len(docs))
# Pages in the original document:  22
# Length of chunks after splitting pages:  353

# Load a Notion database
loader = NotionDirectoryLoader("docs/Notion_DB")
notion_db = loader.load()

docs = text_splitter.split_documents(notion_db)

print("Pages in the original notion document: ", len(notion_db))
print("Length of chunks after splitting pages: ", len(docs))
# Pages in the original notion document:  52
# Length of chunks after splitting pages:  353

from langchain.text_splitter import TokenTextSplitter

text_splitter = TokenTextSplitter(chunk_size=1, chunk_overlap=0)
text1 = "foo bar bazzyfoo"
print(text_splitter.split_text(text1))  
# Output: ['foo', ' bar', ' b', 'az', 'zy', 'foo']

text_splitter = TokenTextSplitter(chunk_size=10, chunk_overlap=0)
docs = text_splitter.split_documents(pages)
print(docs[0])  
# Output: Document(page_content='MachineLearning-Lecture01  \n', metadata={'source': 'docs/cs229_lectures/MachineLearning-Lecture01.pdf', 'page': 0})
print(pages[0].metadata)
# Output: {'source': 'docs/cs229_lectures/MachineLearning-Lecture01.pdf', 'page': 0}

from langchain.document_loaders import NotionDirectoryLoader
from langchain.text_splitter import MarkdownHeaderTextSplitter

markdown_document = """# Title\n\n \
## Chapter 1\n\n \
Hi this is Jim\n\n Hi this is Joe\n\n \
### Section \n\n \
Hi this is Lance \n\n 
## Chapter 2\n\n \
Hi this is Molly"""

headers_to_split_on = [
    ("#", "Header 1"),
    ("##", "Header 2"),
    ("###", "Header 3"),
]

markdown_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on)
md_header_splits = markdown_splitter.split_text(markdown_document)

print(md_header_splits[0])  
# Output: Document(page_content='Hi this is Jim  \nHi this is Joe', metadata={'Header 1': 'Title', 'Header 2': 'Chapter 1'})
print(md_header_splits[1])  
# Output: Document(page_content='Hi this is Lance', metadata={'Header 1': 'Title', 'Header 2': 'Chapter 1', 'Header 3': 'Section'})

loader = NotionDirectoryLoader("docs/Notion_DB")
docs = loader.load()
txt = " ".join([d.page_content for d in docs])

headers_to_split_on = [
    ("#", "Header 1"),
    ("##", "Header 2"),
]
markdown_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on)
md_header_splits = markdown_splitter.split_text(txt)

print(md_header_splits[0])