diff --git a/.gitignore b/.gitignore index 21d1e70..2ac1281 100644 --- a/.gitignore +++ b/.gitignore @@ -37,3 +37,6 @@ htmlcov/ /corpus/ harbor_commons.db /audits/ + +# Downloaded external datasets (large — not source) +/data/external/ diff --git a/packages/ask-a-sailor/src/rag/agent.py b/packages/ask-a-sailor/src/rag/agent.py index 56fbcf5..cae1339 100644 --- a/packages/ask-a-sailor/src/rag/agent.py +++ b/packages/ask-a-sailor/src/rag/agent.py @@ -90,11 +90,15 @@ def search( class AskASailorAgent: + # External corpus sources that can be enabled via --sources + EXTERNAL_SOURCES = ("empathetic_dialogues", "social_chemistry") + def __init__( self, corpus_dir: Path, club_filter: Optional[str] = None, model: str = "gpt-4o-mini", + sources: Optional[list[str]] = None, ): self.client = OpenAI() self.model = model @@ -115,6 +119,19 @@ def __init__( else: print(f" ⚠️ No corpus found for {slug} — run ingestion first.") + # Load optional external corpus sources (e.g. empathetic_dialogues, social_chemistry) + for source in sources or []: + source_path = corpus_dir / source / "corpus.jsonl" + embeddings_path = corpus_dir / source / "embeddings.npy" + if source_path.exists(): + print(f" Loading external source: {source}") + store = SimpleVectorStore() + store.load_from_jsonl(source_path, embeddings_path) + self.store.chunks.extend(store.chunks) + self.store.embeddings.extend(store.embeddings) + else: + print(f" ⚠️ No corpus found for source '{source}' — run ingestion first.") + def embed_query(self, query: str) -> list[float]: response = self.client.embeddings.create( model="text-embedding-3-small", @@ -233,13 +250,23 @@ def main(): default="/tmp/full-harbor/corpus", help="Path to corpus directory", ) + parser.add_argument( + "--sources", + default="", + help=( + "Comma-separated optional external corpus sources to load " + "(e.g. empathetic_dialogues,social_chemistry)" + ), + ) parser.add_argument("--verbose", action="store_true") args = parser.parse_args() + sources = [s.strip() for s in args.sources.split(",") if s.strip()] club_filter = None if args.club == "all" else args.club agent = AskASailorAgent( corpus_dir=Path(args.corpus_dir), club_filter=club_filter, + sources=sources, ) if args.question: diff --git a/packages/ask-a-sailor/tests/test_external_corpus.py b/packages/ask-a-sailor/tests/test_external_corpus.py new file mode 100644 index 0000000..8156b3e --- /dev/null +++ b/packages/ask-a-sailor/tests/test_external_corpus.py @@ -0,0 +1,391 @@ +""" +Tests for external corpus ingestion scripts (EmpatheticDialogues + Social Chemistry 101). +Verifies chunk format, PII filtering, and chunk counts without downloading real datasets. +""" + +import csv +import json +import re +import sys +from pathlib import Path + +import pytest + +# Allow importing ingestion scripts +sys.path.insert(0, str(Path(__file__).parents[3] / "scripts")) +# Allow importing package source +sys.path.insert(0, str(Path(__file__).parents[1] / "src")) + +from ingest_empathetic_dialogues import ( + conversation_to_chunk, + ingest as ingest_ed, + parse_conversations, +) +from ingest_social_chemistry import ( + ingest as ingest_sc, + rot_to_chunk, +) + +# --------------------------------------------------------------------------- +# Required chunk keys (same format as the rest of the corpus) +# --------------------------------------------------------------------------- + +REQUIRED_KEYS = { + "chunk_id", + "doc_id", + "club_slug", + "source_type", + "source_url", + "title", + "text", + "metadata", +} + + +# --------------------------------------------------------------------------- +# Fixtures — tiny synthetic datasets written to tmp directories +# --------------------------------------------------------------------------- + +SAMPLE_ED_ROWS = [ + { + "conv_id": "hit:0_conv:1", + "utterance_idx": "1", + "context": "sentimental", + "prompt": "I remember going to the fireworks with my best friend.", + "speaker_idx": "0", + "utterance": "I remember going to the fireworks with my best friend.", + "selfeval": "", + "tags": "", + }, + { + "conv_id": "hit:0_conv:1", + "utterance_idx": "2", + "context": "sentimental", + "prompt": "", + "speaker_idx": "1", + "utterance": "That sounds like a great memory!", + "selfeval": "like", + "tags": "", + }, + { + "conv_id": "hit:0_conv:2", + "utterance_idx": "1", + "context": "afraid", + "prompt": "I heard a strange noise outside at night.", + "speaker_idx": "0", + "utterance": "I heard a strange noise outside at night.", + "selfeval": "", + "tags": "", + }, + { + "conv_id": "hit:0_conv:2", + "utterance_idx": "2", + "context": "afraid", + "prompt": "", + "speaker_idx": "1", + "utterance": "That can be really scary. Did you find out what it was?", + "selfeval": "", + "tags": "", + }, +] + + +@pytest.fixture() +def ed_input_dir(tmp_path): + """Write a tiny EmpatheticDialogues CSV to a temp directory.""" + d = tmp_path / "empatheticdialogues" + d.mkdir() + csv_path = d / "train.csv" + with open(csv_path, "w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter(f, fieldnames=list(SAMPLE_ED_ROWS[0].keys())) + writer.writeheader() + for row in SAMPLE_ED_ROWS: + writer.writerow(row) + return d + + +SAMPLE_SC_ROWS = [ + { + "area": "amitheasshole", + "m": "", + "rot-agree": "4.0", + "rot-categorization": "morality-ethics|social-norms", + "rot-moral-foundations": "", + "rot-char-targeting": "", + "rot-bad": "0", + "rot-judgment": "it's wrong", + "action": "taking someone else's food", + "rot": "It's wrong to take someone else's food without asking.", + "rot-id": "rot_0001", + "split": "train", + "situation": "taking my roommate's leftovers from the fridge", + "situation-short-id": "sit_001", + "worker-id": "w1", + "n-characters": "1", + "characters": "narrator", + }, + { + "area": "amitheasshole", + "m": "", + "rot-agree": "4.5", + "rot-categorization": "social-norms", + "rot-moral-foundations": "", + "rot-char-targeting": "", + "rot-bad": "0", + "rot-judgment": "it's good", + "action": "helping a friend move", + "rot": "It's good to help a friend when they need it.", + "rot-id": "rot_0002", + "split": "train", + "situation": "my friend asked me to help them move apartments", + "situation-short-id": "sit_002", + "worker-id": "w2", + "n-characters": "2", + "characters": "narrator|friend", + }, +] + + +@pytest.fixture() +def sc_input_dir(tmp_path): + """Write a tiny Social Chemistry 101 TSV to a temp directory.""" + d = tmp_path / "social-chemistry-101" + d.mkdir() + tsv_path = d / "social-chem-101.v1.0.tsv" + with open(tsv_path, "w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter( + f, fieldnames=list(SAMPLE_SC_ROWS[0].keys()), delimiter="\t" + ) + writer.writeheader() + for row in SAMPLE_SC_ROWS: + writer.writerow(row) + return d + + +# --------------------------------------------------------------------------- +# EmpatheticDialogues tests +# --------------------------------------------------------------------------- + +class TestEmpatheticDialogues: + + def test_parse_conversations(self, ed_input_dir): + convs = parse_conversations(ed_input_dir) + assert len(convs) == 2 + assert "hit:0_conv:1" in convs + assert len(convs["hit:0_conv:1"]) == 2 + + def test_chunk_format(self, ed_input_dir): + convs = parse_conversations(ed_input_dir) + chunk = conversation_to_chunk("hit:0_conv:1", convs["hit:0_conv:1"]) + assert chunk is not None + assert REQUIRED_KEYS.issubset(chunk.keys()), ( + f"Missing keys: {REQUIRED_KEYS - chunk.keys()}" + ) + + def test_chunk_source_type(self, ed_input_dir): + convs = parse_conversations(ed_input_dir) + chunk = conversation_to_chunk("hit:0_conv:1", convs["hit:0_conv:1"]) + assert chunk["source_type"] == "empathetic_dialogues" + + def test_chunk_has_emotion_context(self, ed_input_dir): + convs = parse_conversations(ed_input_dir) + chunk = conversation_to_chunk("hit:0_conv:1", convs["hit:0_conv:1"]) + assert "sentimental" in chunk["text"].lower() + assert chunk["metadata"]["emotion"] == "sentimental" + + def test_chunk_contains_utterances(self, ed_input_dir): + convs = parse_conversations(ed_input_dir) + chunk = conversation_to_chunk("hit:0_conv:1", convs["hit:0_conv:1"]) + assert "Speaker:" in chunk["text"] + assert "Listener:" in chunk["text"] + assert "fireworks" in chunk["text"] + + def test_no_pii_in_chunks(self, ed_input_dir): + """Verify no phone numbers, emails, or SSNs leak into chunks.""" + chunks = ingest_ed(ed_input_dir, ed_input_dir / "output") + for chunk in chunks: + text = chunk["text"] + assert not re.search(r"\b\d{3}[-.]?\d{3}[-.]?\d{4}\b", text), \ + f"Phone number found: {text}" + assert not re.search( + r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.\w+", text + ), f"Email found: {text}" + assert not re.search(r"\b\d{3}-\d{2}-\d{4}\b", text), \ + f"SSN found: {text}" + + def test_chunk_count(self, ed_input_dir): + chunks = ingest_ed(ed_input_dir, ed_input_dir / "output") + assert len(chunks) == 2, f"Expected 2 chunks, got {len(chunks)}" + + def test_output_jsonl_valid(self, ed_input_dir): + """Verify the JSONL output is valid line-delimited JSON.""" + out_dir = ed_input_dir / "output" + ingest_ed(ed_input_dir, out_dir) + jsonl_path = out_dir / "corpus.jsonl" + assert jsonl_path.exists() + with open(jsonl_path) as f: + lines = [line for line in f if line.strip()] + assert len(lines) == 2 + for line in lines: + parsed = json.loads(line) + assert REQUIRED_KEYS.issubset(parsed.keys()) + + def test_pii_row_filtered(self, tmp_path): + """A conversation containing a phone number is excluded.""" + d = tmp_path / "ed_pii" + d.mkdir() + rows = [ + { + "conv_id": "pii_conv", + "utterance_idx": "1", + "context": "hopeful", + "prompt": "Call me at 555-123-4567 for details.", + "speaker_idx": "0", + "utterance": "Call me at 555-123-4567 for details.", + "selfeval": "", + "tags": "", + }, + ] + csv_path = d / "train.csv" + with open(csv_path, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=list(rows[0].keys())) + writer.writeheader() + for row in rows: + writer.writerow(row) + + chunks = ingest_ed(d, d / "output") + assert len(chunks) == 0, "PII row should have been filtered" + + def test_comma_replacement(self, ed_input_dir): + """Verify _comma_ placeholders are replaced.""" + rows = [ + { + "conv_id": "comma_conv", + "utterance_idx": "1", + "context": "joyful", + "prompt": "I won first_comma_ second_comma_ and third place!", + "speaker_idx": "0", + "utterance": "I won first_comma_ second_comma_ and third place!", + "selfeval": "", + "tags": "", + }, + ] + chunk = conversation_to_chunk("comma_conv", rows) + assert "_comma_" not in chunk["text"] + assert "first, second, and third" in chunk["text"] + + +# --------------------------------------------------------------------------- +# Social Chemistry 101 tests +# --------------------------------------------------------------------------- + +class TestSocialChemistry: + + def test_chunk_format(self): + chunk = rot_to_chunk(SAMPLE_SC_ROWS[0]) + assert chunk is not None + assert REQUIRED_KEYS.issubset(chunk.keys()), ( + f"Missing keys: {REQUIRED_KEYS - chunk.keys()}" + ) + + def test_chunk_source_type(self): + chunk = rot_to_chunk(SAMPLE_SC_ROWS[0]) + assert chunk["source_type"] == "social_chemistry" + + def test_chunk_contains_rot(self): + chunk = rot_to_chunk(SAMPLE_SC_ROWS[0]) + assert "It's wrong to take someone else's food" in chunk["text"] + + def test_chunk_contains_situation(self): + chunk = rot_to_chunk(SAMPLE_SC_ROWS[0]) + assert "roommate" in chunk["text"] + + def test_chunk_contains_judgment(self): + chunk = rot_to_chunk(SAMPLE_SC_ROWS[0]) + assert "it's wrong" in chunk["text"].lower() + assert chunk["metadata"]["judgment"] == "it's wrong" + + def test_no_pii_in_chunks(self, sc_input_dir): + chunks = ingest_sc(sc_input_dir, sc_input_dir / "output") + for chunk in chunks: + text = chunk["text"] + assert not re.search(r"\b\d{3}[-.]?\d{3}[-.]?\d{4}\b", text) + assert not re.search(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.\w+", text) + assert not re.search(r"\b\d{3}-\d{2}-\d{4}\b", text) + + def test_chunk_count(self, sc_input_dir): + chunks = ingest_sc(sc_input_dir, sc_input_dir / "output") + assert len(chunks) == 2, f"Expected 2 chunks, got {len(chunks)}" + + def test_output_jsonl_valid(self, sc_input_dir): + out_dir = sc_input_dir / "output" + ingest_sc(sc_input_dir, out_dir) + jsonl_path = out_dir / "corpus.jsonl" + assert jsonl_path.exists() + with open(jsonl_path) as f: + lines = [line for line in f if line.strip()] + assert len(lines) == 2 + for line in lines: + parsed = json.loads(line) + assert REQUIRED_KEYS.issubset(parsed.keys()) + + def test_empty_rot_filtered(self): + """A row with an empty ROT is skipped.""" + row = dict(SAMPLE_SC_ROWS[0]) + row["rot"] = "" + chunk = rot_to_chunk(row) + assert chunk is None + + def test_pii_row_filtered(self): + """A row containing an email in the ROT is excluded.""" + row = dict(SAMPLE_SC_ROWS[0]) + row["rot"] = "Contact john@example.com for advice." + chunk = rot_to_chunk(row) + assert chunk is None + + +# --------------------------------------------------------------------------- +# Agent --sources integration +# --------------------------------------------------------------------------- + +class TestAgentSourcesFlag: + + def test_agent_loads_external_source(self, tmp_path): + """ + AskASailorAgent loads external corpus when sources are specified. + """ + from rag.agent import SimpleVectorStore + + # Create a fake external corpus + ext_dir = tmp_path / "social_chemistry" + ext_dir.mkdir() + chunk = { + "chunk_id": "test_c0", + "doc_id": "test_doc", + "club_slug": "", + "source_type": "social_chemistry", + "source_url": "allenai/social-chemistry-101", + "title": "Social norm — test", + "text": "It is good to share.", + "metadata": {"dataset": "Social Chemistry 101"}, + } + with open(ext_dir / "corpus.jsonl", "w") as f: + f.write(json.dumps(chunk) + "\n") + + store = SimpleVectorStore() + store.load_from_jsonl(ext_dir / "corpus.jsonl") + assert len(store.chunks) == 1 + assert store.chunks[0]["source_type"] == "social_chemistry" + + def test_agent_sources_missing_is_graceful(self, tmp_path, capsys): + """ + Specifying a source that doesn't exist should warn, not crash. + """ + from rag.agent import SimpleVectorStore + + # No corpus file exists for this source + source_path = tmp_path / "nonexistent_source" / "corpus.jsonl" + assert not source_path.exists() + # The SimpleVectorStore raises FileNotFoundError; the agent + # handles this by printing a warning. We just confirm + # the path does not exist — the agent's __init__ does the check. diff --git a/scripts/ingest_empathetic_dialogues.py b/scripts/ingest_empathetic_dialogues.py new file mode 100644 index 0000000..d4904a9 --- /dev/null +++ b/scripts/ingest_empathetic_dialogues.py @@ -0,0 +1,195 @@ +""" +Ingest EmpatheticDialogues → Ask a Sailor corpus JSONL +====================================================== +Converts Facebook's EmpatheticDialogues dataset into the JSONL chunk +format expected by the Ask a Sailor RAG pipeline. + +Each conversation is grouped by ``conv_id`` and emitted as a single +chunk whose text includes the emotional context, the situation prompt, +and all utterances labelled by speaker role. + +Dataset: https://dl.fbaipublicfiles.com/parlai/empatheticdialogues/empatheticdialogues.tar.gz +License: CC BY 4.0 + +Usage +----- +1. Download and extract to ``data/external/empatheticdialogues/``:: + + mkdir -p data/external + wget -qO- https://dl.fbaipublicfiles.com/parlai/empatheticdialogues/empatheticdialogues.tar.gz | tar xz -C data/external/ + +2. Run this script:: + + python scripts/ingest_empathetic_dialogues.py \\ + --input-dir data/external/empatheticdialogues \\ + --output-dir /tmp/full-harbor/corpus/empathetic_dialogues +""" + +from __future__ import annotations + +import argparse +import csv +import hashlib +import json +import logging +import re +from collections import defaultdict +from pathlib import Path + +logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") +logger = logging.getLogger(__name__) + +# Columns in the EmpatheticDialogues CSV files +EXPECTED_COLUMNS = [ + "conv_id", "utterance_idx", "context", "prompt", + "speaker_idx", "utterance", "selfeval", "tags", +] + +# Simple patterns that indicate PII leakage (phone, email, SSN) +PII_PATTERNS = [ + re.compile(r"\b\d{3}[-.]?\d{3}[-.]?\d{4}\b"), # phone + re.compile(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.\w+"), # email + re.compile(r"\b\d{3}-\d{2}-\d{4}\b"), # SSN +] + + +def _has_pii(text: str) -> bool: + """Return True if *text* matches any PII pattern.""" + return any(p.search(text) for p in PII_PATTERNS) + + +def _make_chunk_id(text: str) -> str: + return hashlib.sha256(text.encode()).hexdigest()[:12] + + +def parse_conversations(input_dir: Path) -> dict[str, list[dict]]: + """ + Read all EmpatheticDialogues CSV splits and group rows by conv_id. + + Returns a dict mapping ``conv_id`` → list of row dicts, ordered by + ``utterance_idx``. + """ + conversations: dict[str, list[dict]] = defaultdict(list) + + for split in ("train", "valid", "test"): + path = input_dir / f"{split}.csv" + if not path.exists(): + logger.warning("Split file not found: %s", path) + continue + + with open(path, newline="", encoding="utf-8") as f: + reader = csv.DictReader(f) + for row in reader: + conv_id = row.get("conv_id", "").strip() + if not conv_id: + continue + conversations[conv_id].append(row) + + # Sort utterances within each conversation + for conv_id in conversations: + conversations[conv_id].sort( + key=lambda r: int(r.get("utterance_idx", 0)) + ) + return dict(conversations) + + +def conversation_to_chunk(conv_id: str, rows: list[dict]) -> dict | None: + """ + Convert a single conversation (list of rows) into a corpus chunk dict. + + Returns ``None`` if the conversation is empty or contains PII. + """ + if not rows: + return None + + context = rows[0].get("context", "unknown").strip() + prompt = rows[0].get("prompt", "").strip().replace("_comma_", ",") + + lines: list[str] = [] + if prompt: + lines.append(f"Situation: {prompt}") + lines.append(f"Emotion: {context}") + lines.append("") + + for row in rows: + speaker = "Speaker" if row.get("speaker_idx", "0") == "0" else "Listener" + utterance = row.get("utterance", "").strip() + if utterance: + utterance = utterance.replace("_comma_", ",") + lines.append(f"{speaker}: {utterance}") + + text = "\n".join(lines).strip() + if not text or _has_pii(text): + return None + + doc_id = _make_chunk_id(conv_id) + return { + "chunk_id": f"{doc_id}_c0", + "doc_id": doc_id, + "club_slug": "", + "source_type": "empathetic_dialogues", + "source_url": "facebook/empathetic_dialogues", + "title": f"Empathetic dialogue — {context}", + "text": text, + "metadata": { + "emotion": context, + "dataset": "EmpatheticDialogues", + "license": "CC BY 4.0", + }, + } + + +def ingest(input_dir: Path, output_dir: Path) -> list[dict]: + """ + Full pipeline: parse → convert → filter → write JSONL. + + Returns the list of emitted chunks. + """ + conversations = parse_conversations(input_dir) + logger.info("Parsed %d conversations from %s", len(conversations), input_dir) + + chunks: list[dict] = [] + skipped = 0 + for conv_id, rows in conversations.items(): + chunk = conversation_to_chunk(conv_id, rows) + if chunk is None: + skipped += 1 + continue + chunks.append(chunk) + + logger.info( + "Converted %d chunks (%d skipped for PII or empty)", + len(chunks), + skipped, + ) + + output_dir.mkdir(parents=True, exist_ok=True) + out_path = output_dir / "corpus.jsonl" + with open(out_path, "w", encoding="utf-8") as f: + for chunk in chunks: + f.write(json.dumps(chunk, ensure_ascii=False) + "\n") + + logger.info("Saved %d chunks → %s", len(chunks), out_path) + return chunks + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Ingest EmpatheticDialogues into Ask a Sailor corpus JSONL" + ) + parser.add_argument( + "--input-dir", + default="data/external/empatheticdialogues", + help="Path to extracted EmpatheticDialogues directory", + ) + parser.add_argument( + "--output-dir", + default="/tmp/full-harbor/corpus/empathetic_dialogues", + help="Where to write corpus.jsonl", + ) + args = parser.parse_args() + ingest(Path(args.input_dir), Path(args.output_dir)) + + +if __name__ == "__main__": + main() diff --git a/scripts/ingest_social_chemistry.py b/scripts/ingest_social_chemistry.py new file mode 100644 index 0000000..af1f4b5 --- /dev/null +++ b/scripts/ingest_social_chemistry.py @@ -0,0 +1,176 @@ +""" +Ingest Social Chemistry 101 → Ask a Sailor corpus JSONL +======================================================== +Converts the Social Chemistry 101 dataset (Allen AI) into the JSONL +chunk format expected by the Ask a Sailor RAG pipeline. + +Each *rule-of-thumb* (ROT) is emitted as a single chunk containing +the social situation, the rule, and the moral judgment. + +Dataset: https://github.com/mbforbes/social-chemistry-101 +Download: https://storage.googleapis.com/ai2-mosaic-public/projects/social-chemistry/data/social-chem-101.zip +License: CC BY-SA 4.0 + +Usage +----- +1. Download and extract to ``data/external/social-chemistry-101/``:: + + mkdir -p data/external + wget -qO /tmp/social-chem-101.zip \\ + https://storage.googleapis.com/ai2-mosaic-public/projects/social-chemistry/data/social-chem-101.zip + unzip /tmp/social-chem-101.zip -d data/external/social-chemistry-101/ + +2. Run this script:: + + python scripts/ingest_social_chemistry.py \\ + --input-dir data/external/social-chemistry-101 \\ + --output-dir /tmp/full-harbor/corpus/social_chemistry +""" + +from __future__ import annotations + +import argparse +import csv +import hashlib +import json +import logging +import re +from pathlib import Path + +logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") +logger = logging.getLogger(__name__) + +# Simple patterns that indicate PII leakage (phone, email, SSN) +PII_PATTERNS = [ + re.compile(r"\b\d{3}[-.]?\d{3}[-.]?\d{4}\b"), # phone + re.compile(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.\w+"), # email + re.compile(r"\b\d{3}-\d{2}-\d{4}\b"), # SSN +] + + +def _has_pii(text: str) -> bool: + """Return True if *text* matches any PII pattern.""" + return any(p.search(text) for p in PII_PATTERNS) + + +def _make_chunk_id(text: str) -> str: + return hashlib.sha256(text.encode()).hexdigest()[:12] + + +def find_tsv(input_dir: Path) -> Path | None: + """Locate the main TSV file inside *input_dir*.""" + candidates = list(input_dir.rglob("social-chem-101*.tsv")) + if candidates: + return candidates[0] + # Fallback: any .tsv + candidates = list(input_dir.rglob("*.tsv")) + return candidates[0] if candidates else None + + +def rot_to_chunk(row: dict) -> dict | None: + """ + Convert a single ROT row into a corpus chunk dict. + + Returns ``None`` if the row is empty or contains PII. + """ + rot = row.get("rot", "").strip() + situation = row.get("situation", "").strip() + action = row.get("action", "").strip() + judgment = row.get("rot-judgment", "").strip() + rot_id = row.get("rot-id", "").strip() + + if not rot: + return None + + lines: list[str] = [] + if situation: + lines.append(f"Situation: {situation}") + if action: + lines.append(f"Action: {action}") + lines.append(f"Rule of thumb: {rot}") + if judgment: + lines.append(f"Judgment: {judgment}") + + text = "\n".join(lines).strip() + if _has_pii(text): + return None + + doc_id = _make_chunk_id(rot_id or text) + return { + "chunk_id": f"{doc_id}_c0", + "doc_id": doc_id, + "club_slug": "", + "source_type": "social_chemistry", + "source_url": "allenai/social-chemistry-101", + "title": f"Social norm — {rot[:80]}", + "text": text, + "metadata": { + "judgment": judgment, + "area": row.get("area", "").strip(), + "dataset": "Social Chemistry 101", + "license": "CC BY-SA 4.0", + }, + } + + +def ingest(input_dir: Path, output_dir: Path) -> list[dict]: + """ + Full pipeline: locate TSV → parse → convert → filter → write JSONL. + + Returns the list of emitted chunks. + """ + tsv_path = find_tsv(input_dir) + if tsv_path is None: + logger.error("No TSV file found in %s", input_dir) + return [] + + logger.info("Reading %s", tsv_path) + + chunks: list[dict] = [] + skipped = 0 + + with open(tsv_path, newline="", encoding="utf-8") as f: + reader = csv.DictReader(f, delimiter="\t") + for row in reader: + chunk = rot_to_chunk(row) + if chunk is None: + skipped += 1 + continue + chunks.append(chunk) + + logger.info( + "Converted %d chunks (%d skipped for PII or empty)", + len(chunks), + skipped, + ) + + output_dir.mkdir(parents=True, exist_ok=True) + out_path = output_dir / "corpus.jsonl" + with open(out_path, "w", encoding="utf-8") as f: + for chunk in chunks: + f.write(json.dumps(chunk, ensure_ascii=False) + "\n") + + logger.info("Saved %d chunks → %s", len(chunks), out_path) + return chunks + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Ingest Social Chemistry 101 into Ask a Sailor corpus JSONL" + ) + parser.add_argument( + "--input-dir", + default="data/external/social-chemistry-101", + help="Path to extracted Social Chemistry 101 directory", + ) + parser.add_argument( + "--output-dir", + default="/tmp/full-harbor/corpus/social_chemistry", + help="Where to write corpus.jsonl", + ) + args = parser.parse_args() + ingest(Path(args.input_dir), Path(args.output_dir)) + + +if __name__ == "__main__": + main()