scripts/augment_dataset_text.py

from __future__ import annotations

import argparse
import itertools
import json
import random
from pathlib import Path
from typing import Any

from sentence_transformers import SentenceTransformer


def _classify_genre(text: str) -> str | None:
    text_lower = text.lower()
    genre_keywords = {
        "citrus_cologne": ["citrus", "lemon", "lime", "bergamot", "orange", "grapefruit", "cologne", "fresh", "zesty"],
        "floral_woody": ["rose", "jasmine", "floral", "lily", "ylang", "neroli", "flower", "wood", "sandalwood", "cedar"],
        "fougere": ["fougere", "lavender", "coumarin", "oakmoss", "herbal", "green", "moss"],
        "amber_oriental": ["amber", "oriental", "vanilla", "spice", "cinnamon", "incense", "resin", "balsamic", "sweet"],
    }
    scores = {genre: sum(1 for kw in keywords if kw in text_lower) for genre, keywords in genre_keywords.items()}
    best = max(scores.items(), key=lambda item: item[1])
    return best[0] if best[1] > 0 else None


def _load_poucher_profiles(path: Path) -> dict[str, list[str]]:
    """Load Poucher structured profiles into a genre-keyed description pool."""
    descriptions: dict[str, list[str]] = {}
    if not path.exists():
        return descriptions

    for line in path.read_text().strip().splitlines():
        rec = json.loads(line)
        text = rec.get("profile_text", "")
        if not text:
            continue
        genre = _classify_genre(text)
        if genre:
            descriptions.setdefault(genre, []).append(text)
    return descriptions


def load_descriptions_by_genre(
    literature_path: Path,
    arctander_path: Path | None,
    poucher_path: Path | None,
) -> dict[str, list[str]]:
    """Build a genre-keyed pool of perfume text descriptions from all sources."""
    descriptions: dict[str, list[str]] = {}
    genre_keywords = {
        "citrus_cologne": ["citrus", "lemon", "lime", "bergamot", "orange", "grapefruit", "cologne", "fresh", "zesty"],
        "floral_woody": ["rose", "jasmine", "floral", "lily", "ylang", "neroli", "flower", "wood", "sandalwood", "cedar"],
        "fougere": ["fougere", "lavender", "coumarin", "oakmoss", "herbal", "green", "moss"],
        "amber_oriental": ["amber", "oriental", "vanilla", "spice", "cinnamon", "incense", "resin", "balsamic", "sweet"],
    }

    def classify(text: str) -> str | None:
        text_lower = text.lower()
        scores = {genre: sum(1 for kw in keywords if kw in text_lower) for genre, keywords in genre_keywords.items()}
        best = max(scores, key=scores.get)
        return best if scores[best] > 0 else None

    # Collect descriptions from literature formulas
    if literature_path.exists():
        data = json.loads(literature_path.read_text())
        for rec in data:
            genre = rec.get("expected_profile")
            text = rec.get("name", "")
            if genre and text:
                descriptions.setdefault(genre, []).append(text)

    # Collect descriptions from Arctander monographs
    if arctander_path and arctander_path.exists():
        for line in arctander_path.read_text().strip().splitlines():
            rec = json.loads(line)
            text = rec.get("odor_description") or rec.get("raw_text", "")
            if not text:
                continue
            genre = classify(text)
            if genre:
                descriptions.setdefault(genre, []).append(text)

    # Collect structured Poucher profiles
    if poucher_path and poucher_path.exists():
        poucher = _load_poucher_profiles(poucher_path)
        for genre, pool in poucher.items():
            descriptions.setdefault(genre, []).extend(pool)

    # Ensure every genre has a fallback pool
    for genre in genre_keywords:
        descriptions.setdefault(genre, [f"A classic {genre.replace('_', ' ')} fragrance"])

    return descriptions


def augment_dataset(
    input_path: Path,
    output_path: Path,
    descriptions: dict[str, list[str]],
    seed: int = 42,
    poucher_weight: float = 0.5,
) -> None:
    """Attach a sampled text description and its 384-D embedding to every record.

    poucher_weight controls the fraction of records that receive a structured
    Poucher-style `top: ...; middle: ...; base: ...` string. The remainder are
    filled from the full genre pool (Arctander, literature formulas, etc.).
    """
    random.seed(seed)

    print("Loading sentence transformer...")
    text_encoder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

    # Separate structured Poucher strings from the broader prose pool
    poucher_pool: dict[str, list[str]] = {}
    prose_pool: dict[str, list[str]] = {}

    all_structured: list[str] = []
    for genre, pool in descriptions.items():
        poucher_strings = [t for t in pool if "top:" in t and "middle:" in t and "base:" in t]
        prose_strings = [t for t in pool if t not in poucher_strings]
        if poucher_strings:
            poucher_pool[genre] = poucher_strings[:]
            all_structured.extend(poucher_strings)
        prose_pool[genre] = prose_strings if prose_strings else [f"A classic {genre.replace('_', ' ')} fragrance"]

    # Fallback: genres without their own structured profiles can use any structured profile.
    # This guarantees every genre sees the volatility-aware format during training.
    all_structured = list(set(all_structured))
    for genre in descriptions:
        if genre not in poucher_pool and all_structured:
            poucher_pool[genre] = all_structured[:]

    # Pre-shuffle cyclical iterators
    iterators: dict[str, tuple[itertools.cycle, itertools.cycle]] = {}
    for genre in descriptions:
        p_shuffled = poucher_pool.get(genre, prose_pool[genre])[:]
        c_shuffled = prose_pool[genre][:]
        random.shuffle(p_shuffled)
        random.shuffle(c_shuffled)
        iterators[genre] = (itertools.cycle(p_shuffled), itertools.cycle(c_shuffled))

    output_path.parent.mkdir(parents=True, exist_ok=True)
    total = 0
    poucher_used = 0

    with input_path.open("r", encoding="utf-8") as fin, output_path.open("w", encoding="utf-8") as fout:
        for line in fin:
            record = json.loads(line)
            genre = record.get("genre")
            if not genre or genre not in iterators:
                genre = "citrus_cologne"  # fallback

            p_iter, c_iter = iterators[genre]
            if random.random() < poucher_weight and p_iter is not c_iter:
                text = next(p_iter)
                poucher_used += 1
            else:
                text = next(c_iter)

            embedding = text_encoder.encode(text, convert_to_numpy=True)

            record["text_conditioning"] = text
            record["text_embedding"] = embedding.tolist()
            fout.write(json.dumps(record, ensure_ascii=False) + "\n")
            total += 1

    print(f"Augmented {total} records with text conditioning -> {output_path}")
    print(f"Poucher-style strings used: {poucher_used} ({100 * poucher_used / total:.1f}%)")


def main() -> int:
    parser = argparse.ArgumentParser(description="Augment PINO synthetic dataset with text conditioning")
    parser.add_argument("--input", default="data/synthetic_dataset_v2.jsonl", help="Input JSONL")
    parser.add_argument("--output", default="data/synthetic_dataset_v2_text.jsonl", help="Output JSONL")
    parser.add_argument("--literature", default="data/literature_formulas.json", help="Literature blueprint")
    parser.add_argument("--arctander", default="/home/hermes/fragrance-research/extracted/arctander_monographs.jsonl", help="Arctander monographs")
    parser.add_argument("--poucher-profiles", default="data/literature_profiles_poucher.jsonl", help="Poucher structured profiles")
    parser.add_argument("--poucher-weight", type=float, default=0.5, help="Fraction of records that receive structured Poucher-style text")
    parser.add_argument("--seed", type=int, default=42, help="Random seed")
    args = parser.parse_args()

    descriptions = load_descriptions_by_genre(
        Path(args.literature),
        Path(args.arctander) if args.arctander else None,
        Path(args.poucher_profiles) if args.poucher_profiles else None,
    )
    for genre, pool in descriptions.items():
        poucher_count = sum(1 for t in pool if "top:" in t and "middle:" in t and "base:" in t)
        print(f"  {genre}: {len(pool)} descriptions ({poucher_count} structured)")

    augment_dataset(Path(args.input), Path(args.output), descriptions, seed=args.seed, poucher_weight=args.poucher_weight)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())