Skip to content

Chunkers API

This module provides classes for splitting documents into manageable units for processing by LLMs and for providing context to those units.

Unit Chunkers

Unit chunkers determine how a document is divided into smaller pieces for frame extraction. Each piece is a FrameExtractionUnit.

llm_ie.chunkers.UnitChunker 

UnitChunker()

Bases: ABC

This is the abstract class for frame extraction unit chunker. It chunks a document into units (e.g., sentences). LLMs process unit by unit.

Source code in package/llm-ie/src/llm_ie/chunkers.py 
12
13
14
15
16
17
def __init__(self):
    """
    This is the abstract class for frame extraction unit chunker.
    It chunks a document into units (e.g., sentences). LLMs process unit by unit. 
    """
    pass

chunk abstractmethod 

chunk(
    text: str, doc_id: str = None
) -> List[FrameExtractionUnit]
Parameters:

text : str The document text.

Source code in package/llm-ie/src/llm_ie/chunkers.py 
19
20
21
22
23
24
25
26
27
@abc.abstractmethod
def chunk(self, text:str, doc_id:str=None) -> List[FrameExtractionUnit]:
    """
    Parameters:
    ----------
    text : str
        The document text.
    """
    return NotImplemented

chunk_async async 

chunk_async(
    text: str, doc_id: str = None, executor=None
) -> List[FrameExtractionUnit]

asynchronous version of chunk method.

Source code in package/llm-ie/src/llm_ie/chunkers.py 
29
30
31
32
33
34
async def chunk_async(self, text:str, doc_id:str=None, executor=None) -> List[FrameExtractionUnit]:
    """
    asynchronous version of chunk method.
    """
    loop = asyncio.get_running_loop()
    return await loop.run_in_executor(executor, self.chunk, text, doc_id)

llm_ie.chunkers.WholeDocumentUnitChunker 

WholeDocumentUnitChunker()

Bases: UnitChunker

This class chunks the whole document into a single unit (no chunking).

Source code in package/llm-ie/src/llm_ie/chunkers.py 
37
38
39
40
41
def __init__(self):
    """
    This class chunks the whole document into a single unit (no chunking).
    """
    super().__init__()

chunk 

chunk(
    text: str, doc_id: str = None
) -> List[FrameExtractionUnit]
Parameters:

text : str The document text.

Source code in package/llm-ie/src/llm_ie/chunkers.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
def chunk(self, text:str, doc_id:str=None) -> List[FrameExtractionUnit]:
    """
    Parameters:
    ----------
    text : str
        The document text.
    """
    return [FrameExtractionUnit(
        doc_id=doc_id if doc_id is not None else str(uuid.uuid4()),
        start=0,
        end=len(text),
        text=text
    )]

llm_ie.chunkers.SentenceUnitChunker 

SentenceUnitChunker()

Bases: UnitChunker

This class uses the NLTK PunktSentenceTokenizer to chunk a document into sentences.

Source code in package/llm-ie/src/llm_ie/chunkers.py 
 98
 99
100
101
102
def __init__(self):
    """
    This class uses the NLTK PunktSentenceTokenizer to chunk a document into sentences.
    """
    super().__init__()

chunk 

chunk(
    text: str, doc_id: str = None
) -> List[FrameExtractionUnit]
Parameters:

text : str The document text.

Source code in package/llm-ie/src/llm_ie/chunkers.py 
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
def chunk(self, text:str, doc_id:str=None) -> List[FrameExtractionUnit]:
    """
    Parameters:
    ----------
    text : str
        The document text.
    """
    doc_id = doc_id if doc_id is not None else str(uuid.uuid4())
    sentences = []
    for start, end in self.PunktSentenceTokenizer().span_tokenize(text):
        sentences.append(FrameExtractionUnit(
            doc_id=doc_id,
            start=start,
            end=end,
            text=text[start:end]
        ))    
    return sentences

llm_ie.chunkers.SeparatorUnitChunker 

SeparatorUnitChunker(sep: str)

Bases: UnitChunker

This class chunks a document by separator provided.

Parameters:

sep : str a separator string.

Source code in package/llm-ie/src/llm_ie/chunkers.py 
58
59
60
61
62
63
64
65
66
67
68
69
70
71
def __init__(self, sep:str):
    """
    This class chunks a document by separator provided.

    Parameters:
    ----------
    sep : str
        a separator string.
    """
    super().__init__()
    if not isinstance(sep, str):
        raise ValueError("sep must be a string")

    self.sep = sep

chunk 

chunk(
    text: str, doc_id: str = None
) -> List[FrameExtractionUnit]
Parameters:

text : str The document text.

Source code in package/llm-ie/src/llm_ie/chunkers.py 
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
def chunk(self, text:str, doc_id:str=None) -> List[FrameExtractionUnit]:
    """
    Parameters:
    ----------
    text : str
        The document text.
    """
    doc_id = doc_id if doc_id is not None else str(uuid.uuid4())
    paragraphs = text.split(self.sep)
    paragraph_units = []
    start = 0
    for paragraph in paragraphs:
        end = start + len(paragraph)
        paragraph_units.append(FrameExtractionUnit(
            doc_id=doc_id,
            start=start,
            end=end,
            text=paragraph
        ))
        start = end + len(self.sep)
    return paragraph_units

llm_ie.chunkers.TextLineUnitChunker 

TextLineUnitChunker()

Bases: UnitChunker

This class chunks a document into lines.

Source code in package/llm-ie/src/llm_ie/chunkers.py 
124
125
126
127
128
def __init__(self):
    """
    This class chunks a document into lines.
    """
    super().__init__()

chunk 

chunk(
    text: str, doc_id: str = None
) -> List[FrameExtractionUnit]
Parameters:

text : str The document text.

Source code in package/llm-ie/src/llm_ie/chunkers.py 
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
def chunk(self, text:str, doc_id:str=None) -> List[FrameExtractionUnit]:
    """
    Parameters:
    ----------
    text : str
        The document text.
    """
    doc_id = doc_id if doc_id is not None else str(uuid.uuid4())
    lines = text.split('\n')
    line_units = []
    start = 0
    for line in lines:
        end = start + len(line)
        line_units.append(FrameExtractionUnit(
            doc_id=doc_id,
            start=start,
            end=end,
            text=line
        ))
        start = end + 1 
    return line_units

llm_ie.chunkers.LLMUnitChunker 

LLMUnitChunker(
    inference_engine: InferenceEngine,
    prompt_template: str = None,
    system_prompt: str = None,
)

Bases: UnitChunker

This class prompt an LLM for document segmentation (e.g., sections, paragraphs).

Parameters:

inference_engine : InferenceEngine the LLM inferencing engine object. prompt_template : str the prompt template that defines how to chunk the document. Must define a JSON schema with 

[
    {
        "title": "<your title here>",
        "anchor_text": "<the anchor text of the chunk here>"
    },
    {
        "title": "<your title here>",
        "anchor_text": "<the anchor text of the chunk here>"
    }
]
system_prompt : str, optional The system prompt.

Source code in package/llm-ie/src/llm_ie/chunkers.py 
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
def __init__(self, inference_engine:InferenceEngine, prompt_template:str=None, system_prompt:str=None):
    """
    This class prompt an LLM for document segmentation (e.g., sections, paragraphs).

    Parameters:
    ----------
    inference_engine : InferenceEngine
        the LLM inferencing engine object.
    prompt_template : str
        the prompt template that defines how to chunk the document. Must define a JSON schema with 
        ```json
        [
            {
                "title": "<your title here>",
                "anchor_text": "<the anchor text of the chunk here>"
            },
            {
                "title": "<your title here>",
                "anchor_text": "<the anchor text of the chunk here>"
            }
        ]
        ```
    system_prompt : str, optional
        The system prompt.
    """
    self.inference_engine = inference_engine

    if prompt_template is None:
        file_path = importlib.resources.files('llm_ie.asset.default_prompts').joinpath("LLMUnitChunker_user_prompt.txt")
        with open(file_path, 'r', encoding="utf-8") as f:
            self.prompt_template = f.read()
    else:
        self.prompt_template = prompt_template

    self.system_prompt = system_prompt

chunk 

chunk(text, doc_id=None) -> List[FrameExtractionUnit]
Parameters:

text : str the document text. doc_id : str, optional the document id.

Source code in package/llm-ie/src/llm_ie/chunkers.py 
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
def chunk(self, text, doc_id=None) -> List[FrameExtractionUnit]:
    """
    Parameters:
    -----------
    text : str
        the document text.
    doc_id : str, optional
        the document id.
    """
    doc_id = doc_id if doc_id is not None else str(uuid.uuid4())
    user_prompt = apply_prompt_template(prompt_template=self.prompt_template, text_content=text)
    messages = []
    if self.system_prompt is not None:
        messages.append({'role': 'system', 'content': self.system_prompt})
    messages.append({'role': 'user', 'content': user_prompt})

    gen_text = self.inference_engine.chat(messages=messages)

    header_list = extract_json(gen_text=gen_text["response"])
    units = []
    start = 0
    prev_end = 0
    for header in header_list:
        if "anchor_text" not in header:
            Warning.warn(f"Missing anchor_text in header: {header}. Skipping this header.")
            continue
        if not isinstance(header["anchor_text"], str):
            Warning.warn(f"Invalid anchor_text: {header['anchor_text']}. Skipping this header.")
            continue

        start = prev_end
        # find the first instandce of the leading sentence in the rest of the text
        end = text.find(header["anchor_text"], start)
        # if not found, skip this header
        if end == -1:
            continue
        # if start == end (empty text), skip this header
        if start == end:
            continue
        # create a frame extraction unit
        units.append(FrameExtractionUnit(
            doc_id=doc_id,
            start=start,
            end=end,
            text=text[start:end]
        ))
        prev_end = end
    # add the last section
    if prev_end < len(text):
        units.append(FrameExtractionUnit(
            doc_id=doc_id,
            start=prev_end,
            end=len(text),
            text=text[prev_end:]
        ))
    return units

Context Chunkers

Context chunkers determine what contextual information is provided to the LLM alongside a specific FrameExtractionUnit.

llm_ie.chunkers.ContextChunker 

ContextChunker()

Bases: ABC

This is the abstract class for context chunker. Given a frame extraction unit, it returns the context for it.

Source code in package/llm-ie/src/llm_ie/chunkers.py 
248
249
250
251
252
253
def __init__(self):
    """
    This is the abstract class for context chunker. Given a frame extraction unit,
    it returns the context for it.
    """
    pass

fit abstractmethod 

fit(text: str, units: List[FrameExtractionUnit])
Parameters:

text : str The document text.

Source code in package/llm-ie/src/llm_ie/chunkers.py 
255
256
257
258
259
260
261
262
263
@abc.abstractmethod
def fit(self, text:str, units:List[FrameExtractionUnit]):
    """
    Parameters:
    ----------
    text : str
        The document text.
    """
    pass

fit_async async 

fit_async(
    text: str,
    units: List[FrameExtractionUnit],
    executor=None,
)

asynchronous version of fit method.

Source code in package/llm-ie/src/llm_ie/chunkers.py 
265
266
267
268
269
270
async def fit_async(self, text:str, units:List[FrameExtractionUnit], executor=None):
    """
    asynchronous version of fit method.
    """
    loop = asyncio.get_running_loop()
    return await loop.run_in_executor(executor, self.fit, text, units)

chunk abstractmethod 

chunk(unit: FrameExtractionUnit) -> str
Parameters:

unit : FrameExtractionUnit The frame extraction unit.

Return : str The context for the frame extraction unit.

Source code in package/llm-ie/src/llm_ie/chunkers.py 
272
273
274
275
276
277
278
279
280
281
282
283
@abc.abstractmethod
def chunk(self, unit:FrameExtractionUnit) -> str:
    """
    Parameters:
    ----------
    unit : FrameExtractionUnit
        The frame extraction unit.

    Return : str 
        The context for the frame extraction unit.
    """
    return NotImplemented

chunk_async async 

chunk_async(
    unit: FrameExtractionUnit, executor=None
) -> str

asynchronous version of chunk method.

Source code in package/llm-ie/src/llm_ie/chunkers.py 
285
286
287
288
289
290
async def chunk_async(self, unit:FrameExtractionUnit, executor=None) -> str:
    """
    asynchronous version of chunk method.
    """
    loop = asyncio.get_running_loop()
    return await loop.run_in_executor(executor, self.chunk, unit)

llm_ie.chunkers.NoContextChunker 

NoContextChunker()

Bases: ContextChunker

This class does not provide any context.

Source code in package/llm-ie/src/llm_ie/chunkers.py 
294
295
296
297
298
def __init__(self):
    """
    This class does not provide any context.
    """
    super().__init__()

fit 

fit(text: str, units: List[FrameExtractionUnit])
Parameters:

text : str The document text.

Source code in package/llm-ie/src/llm_ie/chunkers.py 
300
301
302
303
304
305
306
307
def fit(self, text:str, units:List[FrameExtractionUnit]):
    """
    Parameters:
    ----------
    text : str
        The document text.
    """
    pass

llm_ie.chunkers.WholeDocumentContextChunker 

WholeDocumentContextChunker()

Bases: ContextChunker

This class provides the whole document as context.

Source code in package/llm-ie/src/llm_ie/chunkers.py 
313
314
315
316
317
318
def __init__(self):
    """
    This class provides the whole document as context.
    """
    super().__init__()
    self.text = None

fit 

fit(text: str, units: List[FrameExtractionUnit])
Parameters:

text : str The document text.

Source code in package/llm-ie/src/llm_ie/chunkers.py 
320
321
322
323
324
325
326
327
def fit(self, text:str, units:List[FrameExtractionUnit]):
    """
    Parameters:
    ----------
    text : str
        The document text.
    """
    self.text = text

llm_ie.chunkers.SlideWindowContextChunker 

SlideWindowContextChunker(window_size: int)

Bases: ContextChunker

This class provides a sliding window context. For example, +-2 sentences around a unit sentence.

Source code in package/llm-ie/src/llm_ie/chunkers.py 
335
336
337
338
339
340
341
def __init__(self, window_size:int):
    """
    This class provides a sliding window context. For example, +-2 sentences around a unit sentence. 
    """
    super().__init__()
    self.window_size = window_size
    self.units = None

fit 

fit(text: str, units: List[FrameExtractionUnit])
Parameters:

units : List[FrameExtractionUnit] The list of frame extraction units.

Source code in package/llm-ie/src/llm_ie/chunkers.py 
343
344
345
346
347
348
349
350
def fit(self, text:str, units:List[FrameExtractionUnit]):
    """
    Parameters:
    ----------
    units : List[FrameExtractionUnit]
        The list of frame extraction units.
    """
    self.units = sorted(units)