Spaces:
Running on Zero
Running on Zero
| 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 | |