Vedashriii/cognitive-load-predictor

1"""2src/preprocessor/chunker.py3 4Takes a raw transcript (list of segments) and:5  1. Cleans the text (removes noise like [Music], [Applause])6  2. Groups segments into fixed time windows (default: 30 seconds)7  3. Returns chunks: [{text, start, end}, ...]8 9Why chunk? BERT has a 512-token limit. Chunking also lets us10map difficulty scores back to specific timestamps in the video.11"""12 13import re14import json15import os16 17 18def clean_text(text: str) -> str:19    """20    Remove transcript noise and normalize whitespace.21    """22    # Remove auto-caption artifacts like [Music], [Applause], [Laughter]23    text = re.sub(r"\[.*?\]", "", text)24    # Remove speaker labels like "SPEAKER 1:" or "John:"25    text = re.sub(r"^[A-Z][A-Za-z\s]+:\s*", "", text)26    # Collapse multiple spaces/newlines into one space27    text = re.sub(r"\s+", " ", text)28    return text.strip()29 30 31def chunk_transcript(transcript: list, window_sec: float = 30.0) -> list:32    """33    Group transcript segments into time windows of `window_sec` seconds.34 35    Each chunk = {36        "text"  : combined text of all segments in this window,37        "start" : start time in seconds,38        "end"   : end time in seconds,39        "chunk_id": index40    }41 42    Why 30 seconds? It gives enough text for BERT to work with43    while keeping fine-grained time resolution in the chart.44    """45    if not transcript:46        raise ValueError("Transcript is empty.")47 48    chunks = []49    current_text = ""50    window_start = transcript[0]["start"]51    window_end = window_start52 53    for seg in transcript:54        seg_start = seg["start"]55        seg_end = seg_start + seg.get("duration", 0)56        cleaned = clean_text(seg["text"])57 58        if not cleaned:59            continue60 61        # If this segment starts a new window, save the current chunk62        if seg_start - window_start >= window_sec:63            if current_text.strip():64                chunks.append({65                    "chunk_id": len(chunks),66                    "text": current_text.strip(),67                    "start": round(window_start, 2),68                    "end": round(window_end, 2),69                })70            # Start a new window71            window_start = seg_start72            current_text = ""73 74        current_text += " " + cleaned75        window_end = seg_end76 77    # Don't forget the last chunk78    if current_text.strip():79        chunks.append({80            "chunk_id": len(chunks),81            "text": current_text.strip(),82            "start": round(window_start, 2),83            "end": round(window_end, 2),84        })85 86    # Filter out very short chunks (less than 10 characters)87    chunks = [c for c in chunks if len(c["text"]) >= 10]88 89    print(f"Chunking complete: {len(chunks)} chunks from {len(transcript)} segments "90          f"(window={window_sec}s)")91    return chunks92 93 94def chunk_raw_text(text: str, words_per_chunk: int = 100) -> list:95    """96    For plain text input (not from YouTube).97    Splits text into chunks of roughly `words_per_chunk` words.98    Assigns fake timestamps (60s per chunk) for visualization.99    """100    words = text.split()101    chunks = []102    for i in range(0, len(words), words_per_chunk):103        chunk_words = words[i: i + words_per_chunk]104        chunk_text = " ".join(chunk_words)105        start_time = (i // words_per_chunk) * 60.0106        chunks.append({107            "chunk_id": len(chunks),108            "text": chunk_text,109            "start": start_time,110            "end": start_time + 60.0,111        })112    print(f"Text chunked into {len(chunks)} chunks ({words_per_chunk} words each).")113    return chunks114 115 116def save_chunks(chunks: list, path: str = "data/processed/chunks.json"):117    """Save chunks to JSON."""118    os.makedirs(os.path.dirname(path), exist_ok=True)119    with open(path, "w", encoding="utf-8") as f:120        json.dump(chunks, f, indent=2, ensure_ascii=False)121    print(f"Chunks saved to: {path}")122 123 124def load_chunks(path: str = "data/processed/chunks.json") -> list:125    """Load chunks from JSON."""126    with open(path, "r", encoding="utf-8") as f:127        return json.load(f)128 129 130# ── Quick test ────────────────────────────────────────────────────────────────131if __name__ == "__main__":132    from src.extractor.youtube import load_transcript133 134    transcript = load_transcript("data/raw/transcript.json")135    chunks = chunk_transcript(transcript, window_sec=30)136    save_chunks(chunks)137 138    print("\nFirst 2 chunks:")139    for c in chunks[:2]:140        print(f"  [{c['start']}s → {c['end']}s]  {c['text'][:80]}...")141