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())