File size: 4,351 Bytes
cd70c8b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
"""Utility functions for loading data and common operations."""

import json
import random
from pathlib import Path
from typing import Dict, List

from config import NEUTRAL_PERSONAS, EPISODES_DIR


def load_data_pool(pool_dir: str = "data/neutral_data_pool/generated") -> Dict[str, Dict[str, List[str]]]:
    """Load the neutral data pool into memory.

    Returns:
        {persona: {topic: [text1, text2, ...]}}
    """
    pool_dir = Path(pool_dir)
    pool: Dict[str, Dict[str, List[str]]] = {}

    for persona in NEUTRAL_PERSONAS:
        persona_dir = pool_dir / persona
        if not persona_dir.exists():
            print(f"Warning: Missing data directory for persona '{persona}' at {persona_dir}")
            pool[persona] = {}
            continue

        pool[persona] = {}
        for jsonl_file in persona_dir.glob("*.jsonl"):
            topic = jsonl_file.stem.replace(f"{persona}_", "")
            texts = []
            with open(jsonl_file, "r") as f:
                for line in f:
                    try:
                        record = json.loads(line)
                        texts.append(record.get("text", ""))
                    except json.JSONDecodeError:
                        continue
            if texts:
                pool[persona][topic] = texts

    return pool


def create_minimal_data_pool() -> Dict[str, Dict[str, List[str]]]:
    """Create a minimal in-memory data pool for testing without files."""
    topics = ["politics", "sports", "technology", "health", "entertainment"]
    pool = {}
    for persona in NEUTRAL_PERSONAS:
        pool[persona] = {}
        for topic in topics:
            pool[persona][topic] = [
                f"[{persona}] Sample post about {topic} #{i}"
                for i in range(50)
            ]
    return pool


def load_agenda_config(path: str = "data/episodes/agenda_config.json") -> Dict[str, List[str]]:
    """Load agenda configuration."""
    path = Path(path)
    if not path.exists():
        return {}
    with open(path) as f:
        return json.load(f)


def load_propaganda_bank(bank_dir: str = "data/episodes/propaganda_bank") -> Dict[str, Dict[int, Dict[str, List[str]]]]:
    """Load propaganda bank into memory.

    Returns:
        {topic_slug: {angle_idx: {persona: [text1, text2, ...]}}}
    """
    bank_dir = Path(bank_dir)
    if not bank_dir.exists():
        return {}

    bank: Dict[str, Dict[int, Dict[str, List[str]]]] = {}
    personas = ["hothead", "clickbait", "careful_journalist", "partisan", "rando"]

    for topic_dir in bank_dir.iterdir():
        if not topic_dir.is_dir():
            continue
        topic_slug = topic_dir.name
        bank[topic_slug] = {}

        for angle_dir in topic_dir.iterdir():
            if not angle_dir.is_dir() or not angle_dir.name.startswith("angle_"):
                continue
            angle_idx = int(angle_dir.name.replace("angle_", ""))
            bank[topic_slug][angle_idx] = {}

            for persona_file in angle_dir.glob("*.jsonl"):
                persona = persona_file.stem
                if persona not in personas:
                    continue
                texts = []
                with open(persona_file, "r") as f:
                    for line in f:
                        try:
                            record = json.loads(line)
                            texts.append(record.get("text", ""))
                        except json.JSONDecodeError:
                            continue
                if texts:
                    bank[topic_slug][angle_idx][persona] = texts

    return bank


def get_episode_pool_path() -> Path:
    """Get the directory path for episode pool (relative to project root)."""
    # Find project root by looking for PLANS.md or README.md
    current = Path.cwd()
    while current != current.parent:
        if (current / "PLANS.md").exists() or (current / "README.md").exists():
            path = current / EPISODES_DIR
            path.mkdir(parents=True, exist_ok=True)
            return path
        current = current.parent
    # Fallback to relative path
    path = Path(EPISODES_DIR)
    path.mkdir(parents=True, exist_ok=True)
    return path


def get_all_episode_dirs() -> List[Path]:
    """Get all episode directories."""
    pool_path = get_episode_pool_path()
    return sorted(pool_path.glob("ep_*"))