File size: 3,803 Bytes
e454383 | 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 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | import re
from datetime import datetime
def chunk_text_splitter(documents, separator="\n", chunk_size=1000, chunk_overlap=100):
"""
Splits the 'text' of each document into chunks of specified size, with optional overlap.
Returns a list of dicts with the same metadata as the input.
"""
chunks = []
for doc in documents:
text = doc["text"]
meta = {k: v for k, v in doc.items() if k != "text"}
# Split and join to normalize
split_text = text.split(separator)
joined_text = separator.join(split_text)
start = 0
while start < len(joined_text):
end = min(start + chunk_size, len(joined_text))
chunk = joined_text[start:end]
if chunk.strip():
chunks.append({"text": chunk, **meta})
start += chunk_size - chunk_overlap
return chunks
def paragraph_text_splitter(documents, separator="\n\n"):
"""
Splits the 'text' of each document into paragraphs.
Returns a list of dicts with the same metadata as the input.
"""
chunks = []
for doc in documents:
text = doc["text"]
meta = {k: v for k, v in doc.items() if k != "text"}
for para in text.split(separator):
if para.strip():
chunks.append({"text": para.strip(), **meta})
return chunks
def transcript_chunker(documents):
"""
Splits each transcript document by lines and attaches speaker/timestamp if present.
Returns a list of dicts with the same metadata as the input.
"""
pattern = r"^\[(\d{2}:\d{2}:\d{2})\]\s+([A-Za-z]+):\s*(.*)$"
chunks = []
for doc in documents:
meta = {k: v for k, v in doc.items() if k != "text"}
for line in doc["text"].splitlines():
match = re.match(pattern, line)
if match:
timestamp = match.group(1)
speaker = match.group(2)
content = match.group(3)
timestamp_seconds = sum([int(x) * 60 ** i for i, x in enumerate(reversed(timestamp.split(":")))])
chunk = {
"text": content,
"timestamp": timestamp,
"timestamp_seconds": timestamp_seconds,
"speaker": speaker,
**meta
}
chunks.append(chunk)
elif line.strip():
# If no match, just add the line as text with metadata
chunks.append({"text": line.strip(), **meta})
return chunks
def srt_chunker(documents):
"""
Splits SRT transcript documents into chunks with start/end time, speaker, and content.
Returns a list of dicts with the same metadata as the input.
"""
pattern = r"(\d{2}:\d{2}:\d{2}),\d{3}\s-->\s(\d{2}:\d{2}:\d{2}),\d{3}\n\[(.*?)\]:\s(.+)"
chunks = []
for doc in documents:
meta = {k: v for k, v in doc.items() if k != "text"}
matches = re.finditer(pattern, doc["text"], re.MULTILINE)
for match in matches:
start_time = match.group(1)
end_time = match.group(2)
speaker = match.group(3)
content = match.group(4).replace("\n", " ")
start_time_seconds = sum([int(x) * 60 ** i for i, x in enumerate(reversed(start_time.split(":")))])
end_time_seconds = sum([int(x) * 60 ** i for i, x in enumerate(reversed(end_time.split(":")))])
chunk = {
"text": content,
"start_time": start_time,
"end_time": end_time,
"start_time_seconds": start_time_seconds,
"end_time_seconds": end_time_seconds,
"speaker": speaker,
**meta
}
chunks.append(chunk)
return chunks
|