Replace hardcoded keyword filters with LLM-based target-aware category mapping

The stratified sampler used hardcoded substring matching to bucket 567
occupation values into 7 categories — misclassifying many (e.g. "sales
engineer" → tech) and dumping ~200 categories into "other". Since this is
fundamentally a retrieval problem, the LLM needs to see all categories
relative to the evaluation target.

- Add build_category_map.py: sends unique category values + entity
description to LLM, which returns 6-10 groups based on how those
categories would differ in evaluating the specific entity. Cached per
entity so the LLM call happens once.
- Remove occupation_bucket() and education_tier() from stratified_sampler.py;
education uses raw 7 values directly, occupation uses LLM mapping when
--entity is provided or raw values otherwise.
- Fix generate_cohort.py: add missing fields (sex, country, marital_status,
interests) so generated cohorts are compatible with the rest of the pipeline.
- Remove interests truncation (hobbies[:5] + skills[:3]) in persona_loader.py;
pass through full lists from the dataset.
- Remove unused MBTI_TYPES constant.
- Fix build_category_map.py data path to use project's data/nemotron/
(consistent with setup_data.py) instead of hardcoded ~/data/nvidia/ path.

AGENT.md

@@ -89,9 +89,8 @@ Run the stratified sampler with the user's parameters:
 
 ```bash
 uv run python scripts/stratified_sampler.py \
-  --population <dataset_or_generated> \
-  --filters '{"sex": "Female", "state": "IL", "age_min": 25, "age_max": 50}' \
-  --dimensions '["age_bracket", "marital_status", "education_tier"]' \
+  --input data/filtered.json \
+  --entity entities/<entity>.md \
   --total 50 \
   --output data/cohort.json
 ```

SKILL.md

@@ -71,9 +71,10 @@ uv run python scripts/persona_loader.py \
   --filters '{"sex": "...", "state": "...", "age_min": N, "age_max": N}' \
   --output data/filtered.json
 
-# Stratified sample
+# Stratified sample (--entity enables LLM-based occupation bucketing)
 uv run python scripts/stratified_sampler.py \
   --input data/filtered.json \
+  --entity entities/<entity>.md \
   --total 50 \
   --output data/cohort.json
 ```

scripts/build_category_map.py

@@ -0,0 +1,254 @@
+"""
+Build a target-aware category mapping for stratified sampling.
+
+Instead of hardcoded keyword buckets, this sends all unique category values
+from the dataset to an LLM along with the entity description. The LLM returns
+a grouping that's meaningful for the specific evaluation target.
+
+The mapping is cached so it's only generated once per entity + field combination.
+
+Usage:
+    # Build occupation mapping for a specific entity
+    uv run python scripts/build_category_map.py \
+      --entity entities/my_product.md \
+      --field occupation \
+      --output data/occupation_map.json
+
+    # As a library
+    from build_category_map import load_or_build_map
+    occ_map = load_or_build_map("occupation", entity_text, unique_values)
+"""
+
+import json
+import hashlib
+import os
+import re
+import argparse
+from pathlib import Path
+
+from dotenv import load_dotenv
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+load_dotenv(PROJECT_ROOT / ".env")
+
+from openai import OpenAI
+
+CACHE_DIR = PROJECT_ROOT / "data" / "category_maps"
+
+SYSTEM_PROMPT = """You are a research methodology assistant. Your job is to create
+meaningful groupings of categorical values for stratified sampling.
+
+You will receive:
+1. A list of unique category values from a dataset
+2. A description of the entity being evaluated
+
+Create 6-10 groups that ensure the evaluation cohort captures meaningfully
+different perspectives on the entity. Groups should reflect how people in these
+categories would DIFFER in their evaluation of the entity — not just demographic
+similarity.
+
+You MUST respond with valid JSON only."""
+
+MAP_PROMPT = """## Entity Being Evaluated
+
+{entity}
+
+---
+
+## Unique Values to Group
+
+Field: {field}
+Values ({count} unique):
+
+{values}
+
+---
+
+## Task
+
+Group these {count} values into 6-10 buckets that capture meaningfully different
+perspectives on the entity above. Every value must appear in exactly one bucket.
+
+Think about: Who would evaluate this entity differently? What professional/life
+context changes how someone perceives this?
+
+Return JSON:
+{{
+    "buckets": [
+        {{
+            "name": "<short bucket label>",
+            "rationale": "<why this group evaluates the entity differently>",
+            "values": ["<value1>", "<value2>", ...]
+        }}
+    ]
+}}"""
+
+
+def extract_unique_values(field, data_dir=None):
+    """Extract unique values for a field from the Nemotron dataset.
+
+    Uses HuggingFace datasets (load_from_disk), consistent with setup_data.py
+    and persona_loader.py.
+    """
+    from collections import Counter
+    from datasets import load_from_disk
+
+    if data_dir is None:
+        data_dir = PROJECT_ROOT / "data" / "nemotron"
+
+    if not (data_dir / "dataset_info.json").exists():
+        raise FileNotFoundError(
+            f"Dataset not found at {data_dir}. Run: uv run python scripts/setup_data.py"
+        )
+
+    ds = load_from_disk(str(data_dir))
+    return Counter(ds[field])
+
+
+def build_map(field, entity_text, unique_values, client=None, model=None):
+    """Call LLM to build a target-aware category mapping."""
+    if client is None:
+        client = OpenAI(
+            api_key=os.getenv("LLM_API_KEY"),
+            base_url=os.getenv("LLM_BASE_URL"),
+        )
+    if model is None:
+        model = os.getenv("LLM_MODEL_NAME")
+
+    values_text = "\n".join(f"  - {v}" for v in sorted(unique_values))
+
+    prompt = MAP_PROMPT.format(
+        entity=entity_text,
+        field=field,
+        count=len(unique_values),
+        values=values_text,
+    )
+
+    resp = client.chat.completions.create(
+        model=model,
+        messages=[
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": prompt},
+        ],
+        response_format={"type": "json_object"},
+        max_tokens=16384,
+        temperature=0.3,
+    )
+
+    content = resp.choices[0].message.content
+    if not content:
+        raise RuntimeError("Empty response from LLM")
+    content = re.sub(r'<think>[\s\S]*?</think>', '', content).strip()
+    data = json.loads(content)
+
+    # Flatten to value -> bucket_name mapping
+    mapping = {}
+    for bucket in data["buckets"]:
+        for val in bucket["values"]:
+            mapping[val] = bucket["name"]
+
+    # Check coverage
+    mapped = set(mapping.keys())
+    expected = set(unique_values)
+    missing = expected - mapped
+    if missing:
+        print(f"  Warning: {len(missing)} values not mapped by LLM, assigning to 'other':")
+        for v in sorted(missing)[:10]:
+            print(f"    - {v}")
+        if len(missing) > 10:
+            print(f"    ... and {len(missing) - 10} more")
+        for v in missing:
+            mapping[v] = "other"
+
+    extra = mapped - expected
+    if extra:
+        print(f"  Note: LLM included {len(extra)} values not in dataset (ignored)")
+
+    return {
+        "field": field,
+        "buckets": data["buckets"],
+        "mapping": mapping,
+    }
+
+
+def cache_key(field, entity_text):
+    """Generate a stable cache key from field + entity content."""
+    h = hashlib.sha256(entity_text.encode()).hexdigest()[:12]
+    return f"{field}_{h}"
+
+
+def load_or_build_map(field, entity_text, unique_values,
+                       client=None, model=None, cache_dir=None):
+    """Load cached mapping or build a new one."""
+    cache_dir = Path(cache_dir or CACHE_DIR)
+    cache_dir.mkdir(parents=True, exist_ok=True)
+
+    key = cache_key(field, entity_text)
+    cache_path = cache_dir / f"{key}.json"
+
+    if cache_path.exists():
+        print(f"  Loading cached {field} mapping: {cache_path.name}")
+        with open(cache_path) as f:
+            data = json.load(f)
+        return data["mapping"]
+
+    print(f"  Building {field} mapping ({len(unique_values)} unique values)...")
+    data = build_map(field, entity_text, unique_values, client, model)
+
+    with open(cache_path, "w") as f:
+        json.dump(data, f, ensure_ascii=False, indent=2)
+    print(f"  Cached to {cache_path.name}")
+
+    # Print bucket summary
+    for b in data["buckets"]:
+        print(f"    {b['name']}: {len(b['values'])} values — {b['rationale']}")
+
+    return data["mapping"]
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--entity", required=True, help="Path to entity document")
+    parser.add_argument("--field", default="occupation",
+                        help="Dataset field to map (default: occupation)")
+    parser.add_argument("--data-dir", default=None,
+                        help="Path to Nemotron arrow shards")
+    parser.add_argument("--output", default=None,
+                        help="Output path (default: data/category_maps/<key>.json)")
+    parser.add_argument("--force", action="store_true",
+                        help="Rebuild even if cached")
+    args = parser.parse_args()
+
+    entity_text = Path(args.entity).read_text()
+    counts = extract_unique_values(args.field, args.data_dir and Path(args.data_dir))
+    unique_values = list(counts.keys())
+
+    print(f"Field: {args.field} | {len(unique_values)} unique values")
+    print(f"Entity: {args.entity}")
+
+    if args.force:
+        key = cache_key(args.field, entity_text)
+        cache_path = CACHE_DIR / f"{key}.json"
+        if cache_path.exists():
+            cache_path.unlink()
+            print("  Cleared cache")
+
+    mapping = load_or_build_map(args.field, entity_text, unique_values)
+
+    if args.output:
+        out = Path(args.output)
+        out.parent.mkdir(parents=True, exist_ok=True)
+        with open(out, "w") as f:
+            json.dump(mapping, f, ensure_ascii=False, indent=2)
+        print(f"\nSaved mapping to {args.output}")
+
+    # Summary
+    from collections import Counter
+    bucket_counts = Counter(mapping.values())
+    print(f"\nBucket distribution:")
+    for bucket, cnt in bucket_counts.most_common():
+        print(f"  {bucket}: {cnt} categories")
+
+
+if __name__ == "__main__":
+    main()

scripts/generate_cohort.py

@@ -57,10 +57,14 @@ Return JSON:
         {{
             "name": "<realistic full name>",
             "age": <integer>,
+            "sex": "<Male | Female>",
             "city": "<city>",
             "state": "<state abbreviation>",
+            "country": "USA",
             "education_level": "<high_school | bachelors | graduate | etc>",
             "occupation": "<specific job title>",
+            "marital_status": "<never_married | married | divorced | widowed | separated>",
+            "interests": ["<hobby or skill, 3-5 items>"],
             "persona": "<200-400 word detailed persona narrative>",
             "segment": "{segment_label}"
         }}

scripts/persona_loader.py

@@ -24,11 +24,6 @@ from datasets import load_from_disk
 PROJECT_ROOT = Path(__file__).resolve().parent.parent
 DEFAULT_DATA_DIR = PROJECT_ROOT / "data" / "nemotron"
 
-MBTI_TYPES = [
-    "INTJ", "INTP", "ENTJ", "ENTP", "INFJ", "INFP", "ENFJ", "ENFP",
-    "ISTJ", "ISFJ", "ESTJ", "ESFJ", "ISTP", "ISFP", "ESTP", "ESFP",
-]
-
 # All narrative fields in the dataset, in order of richness
 NARRATIVE_FIELDS = [
     "persona", "cultural_background", "professional_persona",
@@ -150,7 +145,7 @@ def to_profile(row: dict, user_id: int) -> dict:
         "education_level": row.get("education_level", ""),
         "marital_status": row.get("marital_status", ""),
         "occupation": (row.get("occupation") or "").replace("_", " ").title(),
-        "interests": hobbies[:5] + skills[:3],
+        "interests": hobbies + skills,
         "source_uuid": row.get("uuid", ""),
     }
 

scripts/stratified_sampler.py

@@ -5,17 +5,23 @@ Stratification is configurable: pass dimension functions that map a row to a
 bucket label. The sampler ensures minimum 1 per non-empty stratum, then fills
 proportionally with within-stratum diversity on a secondary dimension.
 
+When --entity is provided, occupation bucketing is done via LLM: the full set
+of unique occupation values is sent to the model along with the entity
+description, so the grouping reflects how different professions would evaluate
+that specific entity. The mapping is cached per entity.
+
 Usage:
     uv run python scripts/stratified_sampler.py \
       --input data/filtered.json \
+      --entity entities/my_product.md \
       --total 50 \
       --output data/cohort.json
 
-    # Or with custom dimensions (as Python expressions)
+    # Without entity (uses raw occupation values, no bucketing)
     uv run python scripts/stratified_sampler.py \
       --input data/filtered.json \
       --total 50 \
-      --dim-exprs '["age_bracket(r[\"age\"])", "r[\"marital_status\"]", "education_tier(r[\"education_level\"])"]'
+      --output data/cohort.json
 """
 
 import json
@@ -37,30 +43,40 @@ def age_bracket(age: int) -> str:
     return "50+"
 
 
-def education_tier(edu: str) -> str:
-    if edu in ("graduate",): return "graduate"
-    if edu in ("bachelors",): return "bachelors"
-    if edu in ("associates", "some_college"): return "some_college"
-    return "no_degree"
-
-
-def occupation_bucket(occ: str) -> str:
-    occ = occ.lower()
-    for kw in ("software", "computer", "data", "web", "engineer", "developer"):
-        if kw in occ: return "tech"
-    for kw in ("nurse", "doctor", "physician", "therapist", "health", "medical"):
-        if kw in occ: return "healthcare"
-    for kw in ("teacher", "professor", "instructor", "education"):
-        if kw in occ: return "education"
-    for kw in ("manager", "accountant", "financial", "analyst", "marketing", "sales"):
-        if kw in occ: return "business"
-    for kw in ("artist", "designer", "writer", "musician", "photographer"):
-        if kw in occ: return "creative"
-    for kw in ("cashier", "retail", "food", "customer", "secretary", "laborer"):
-        if kw in occ: return "service"
-    if occ in ("not in workforce", "no occupation", ""):
-        return "not_working"
-    return "other"
+def make_occupation_fn(entity_path=None, profiles=None):
+    """
+    Build an occupation bucketing function.
+
+    With --entity: uses LLM to create a target-aware mapping from the full
+    set of unique occupation values. Cached per entity content.
+
+    Without --entity: passes through the raw occupation value.
+    """
+    if entity_path is None:
+        return lambda p: p.get("occupation", "unknown") or "unknown"
+
+    entity_text = Path(entity_path).read_text()
+
+    # Collect unique occupation values from the profiles being sampled
+    unique_occs = set()
+    for p in (profiles or []):
+        occ = p.get("occupation", "")
+        if occ:
+            unique_occs.add(occ)
+
+    if not unique_occs:
+        return lambda p: p.get("occupation", "unknown") or "unknown"
+
+    import sys
+    sys.path.insert(0, str(Path(__file__).resolve().parent))
+    from build_category_map import load_or_build_map
+    mapping = load_or_build_map("occupation", entity_text, list(unique_occs))
+
+    def lookup(p):
+        occ = p.get("occupation", "")
+        return mapping.get(occ, mapping.get(occ.lower(), "other"))
+
+    return lookup
 
 
 # ── Sampler ───────────────────────────────────────────────────────────────
@@ -141,6 +157,8 @@ def stratified_sample(profiles, dim_fns, total=50, diversity_fn=None, seed=42):
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--input", default="data/filtered.json")
+    parser.add_argument("--entity", default=None,
+                        help="Path to entity document (enables LLM-based occupation bucketing)")
     parser.add_argument("--total", type=int, default=50)
     parser.add_argument("--seed", type=int, default=42)
     parser.add_argument("--output", default="data/cohort.json")
@@ -150,13 +168,16 @@ def main():
         profiles = json.load(f)
     print(f"Loaded {len(profiles)} profiles from {args.input}")
 
-    # Default dimensions: age, marital status, education
+    # Build occupation function — LLM-based if entity provided, raw passthrough otherwise
+    occupation_fn = make_occupation_fn(args.entity, profiles)
+
+    # Default dimensions: age, marital status, education (raw values)
     dim_fns = [
         lambda p: age_bracket(p.get("age", 30)),
         lambda p: p.get("marital_status", "unknown"),
-        lambda p: education_tier(p.get("education_level", "")),
+        lambda p: p.get("education_level", "") or "unknown",
     ]
-    diversity_fn = lambda p: occupation_bucket(p.get("occupation", ""))
+    diversity_fn = occupation_fn
 
     selected = stratified_sample(profiles, dim_fns, total=args.total,
                                   diversity_fn=diversity_fn, seed=args.seed)
@@ -173,8 +194,8 @@ def main():
     print(f"\nSaved {len(selected)} to {args.output}")
     for dim_name, fn in [("Age", lambda p: age_bracket(p.get("age", 30))),
                           ("Marital", lambda p: p.get("marital_status", "?")),
-                          ("Education", lambda p: education_tier(p.get("education_level", ""))),
-                          ("Occupation", lambda p: occupation_bucket(p.get("occupation", "")))]:
+                          ("Education", lambda p: p.get("education_level", "") or "unknown"),
+                          ("Occupation", occupation_fn)]:
         dist = Counter(fn(p) for p in selected)
         print(f"  {dim_name}: {dict(sorted(dist.items()))}")
     print(f"  Cities: {len(set(p.get('city','') for p in selected))} unique")