Spaces:
Running on Zero
Running on Zero
File size: 3,953 Bytes
712c393 | 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 97 98 99 100 101 102 | from __future__ import annotations
import re
from typing import List, Sequence
from ..ingestion.models import DocumentBlock
from .base import BaseChunker
from .models import Chunk
class TranscriptChunker(BaseChunker):
"""
Time-aware chunker for audio and video transcripts.
Preserves start and end timestamps in source locators and metadata.
"""
TIMESTAMP_REGEX = re.compile(r"\[?(\d{1,2}:\d{2}(?::\d{2})?(?:\.\d+)?)\]?")
def chunk(self, blocks: Sequence[DocumentBlock], document_id: str) -> List[Chunk]:
chunks: List[Chunk] = []
target_chars = int(self.config.chunk_size * self.config.chars_per_token)
for block_idx, block in enumerate(blocks):
text = block.text.strip()
if not text:
continue
lines = [line.strip() for line in text.split("\n") if line.strip()]
current_batch: List[str] = []
current_len = 0
start_ts = None
end_ts = None
c_idx = 0
for line in lines:
ts_match = self.TIMESTAMP_REGEX.search(line)
if ts_match:
found_ts = ts_match.group(1)
if start_ts is None:
start_ts = found_ts
end_ts = found_ts
if current_len + len(line) > target_chars and current_batch:
chunk_text = "\n".join(current_batch)
chunk_id = self.make_chunk_id(
document_id, f"transcript_b{block_idx}_c{c_idx}"
)
time_locator = (
f"time:{start_ts or '00:00'}-{end_ts or start_ts or '00:00'}"
)
chunks.append(
Chunk(
chunk_id=chunk_id,
document_id=document_id,
text=chunk_text,
chunk_type="transcript",
source_locator=block.source_locator or time_locator,
token_count=self.estimate_tokens(chunk_text),
metadata={
**block.metadata,
"start_time": start_ts,
"end_time": end_ts,
"is_transcript": True,
},
)
)
c_idx += 1
current_batch = [line]
current_len = len(line)
start_ts = ts_match.group(1) if ts_match else None
else:
current_batch.append(line)
current_len += len(line)
if current_batch:
chunk_text = "\n".join(current_batch)
chunk_id = self.make_chunk_id(
document_id, f"transcript_b{block_idx}_c{c_idx}"
)
time_locator = (
f"time:{start_ts or '00:00'}-{end_ts or start_ts or '00:00'}"
)
chunks.append(
Chunk(
chunk_id=chunk_id,
document_id=document_id,
text=chunk_text,
chunk_type="transcript",
source_locator=block.source_locator or time_locator,
token_count=self.estimate_tokens(chunk_text),
metadata={
**block.metadata,
"start_time": start_ts,
"end_time": end_ts,
"is_transcript": True,
},
)
)
return chunks
|