persona-data 0.1.0__tar.gz

persona_data-0.1.0/PKG-INFO

@@ -0,0 +1,127 @@
+Metadata-Version: 2.3
+Name: persona-data
+Version: 0.1.0
+Summary: Shared dataset loading and prompt formatting for implicit-personalization projects
+Requires-Dist: huggingface-hub>=0.30.0
+Requires-Dist: python-dotenv>=1.0.0
+Requires-Dist: torch>=2.0.0
+Requires-Dist: numpy>=1.24.0
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+
+# persona-data
+
+[![Docs](https://img.shields.io/badge/docs-view-purple?logo=materialformkdocs)](https://implicit-personalization.github.io/persona-data/)
+
+Shared dataset loading, prompt formatting, and environment utilities for the [implicit-personalization](https://github.com/implicit-personalization) projects.
+
+## Overview
+
+`persona-data` provides the common dataset and prompt helpers used across the persona projects:
+
+- `SynthPersonaDataset` for persona profiles plus QA pairs
+- `PersonaGuessDataset` for turn-based persona games
+- prompt helpers for roleplay and multiple-choice evaluation
+- environment helpers for seeds, devices, and artifact paths
+
+## Installation
+
+Add as a uv git source in your project's `pyproject.toml`:
+
+```toml
+[project]
+dependencies = ["persona-data"]
+
+[tool.uv.sources]
+persona-data = { git = "ssh://git@github.com/implicit-personalization/persona-data.git" }
+```
+
+Then run `uv sync`.
+
+For local development alongside other repos, use an editable path source:
+
+```toml
+[tool.uv.sources]
+persona-data = { path = "../persona-data", editable = true }
+```
+
+## Package layout
+
+```
+src/persona_data/
+├── __init__.py
+├── synth_persona.py       # SynthPersonaDataset, PersonaDataset, PersonaData, QAPair, BiographySection
+├── persona_guess.py       # PersonaGuessDataset, GameRecord, Turn
+├── prompts.py             # format_roleplay_prompt, format_mc_question, format_messages
+└── environment.py         # load_env, set_seed, get_device, get_artifacts_dir
+```
+
+## Datasets
+
+Each dataset is a module with its own types and a loader that downloads from Hugging Face, cached via `HF_HOME`.
+
+### SynthPersona
+
+```python
+from persona_data.synth_persona import SynthPersonaDataset
+
+dataset = SynthPersonaDataset()
+
+persona = dataset[0]
+persona.name              # "Ethan Robinson"
+persona.templated_view    # short attribute-based system prompt
+persona.biography_view    # full biography text
+persona.sections          # list of BiographySection
+
+qa_pairs = dataset.get_qa(persona.id, type="implicit", difficulty=[1, 2])
+questions = dataset.questions(persona.id, type="explicit")
+```
+
+### PersonaGuess
+
+```python
+from persona_data.persona_guess import PersonaGuessDataset
+
+games = PersonaGuessDataset()
+game = games[0]
+turns = games.get_qa(game.game_id, player="A")
+questions = games.questions(game.game_id, player="B")
+```
+
+## Prompt formatting
+
+```python
+from persona_data.prompts import format_messages, format_roleplay_prompt
+
+system_prompt = format_roleplay_prompt(persona.biography_view)
+
+messages = [
+    {"role": "system", "content": system_prompt},
+    {"role": "user", "content": "Where did you grow up?"},
+    {"role": "assistant", "content": "I grew up in Little Rock, Arkansas."},
+]
+full_prompt, response_start_idx = format_messages(messages, tokenizer)
+```
+
+`format_roleplay_prompt` supports `mode="roleplay"` (default), `mode="conversational"`, and `mode="mc"`.
+
+`format_messages` handles tokenizers that do not support the `"system"` role (for example Gemma 2) by merging system content into the first user message.
+
+For multiple-choice evaluation, use `format_mc_question(qa)` and `mc_correct_letter(qa)`.
+
+## Environment helpers
+
+```python
+from persona_data.environment import load_env, set_seed, get_device, get_artifacts_dir
+
+load_env()            # loads .env from cwd (searches parent dirs)
+set_seed(1337)        # sets random, numpy, and torch seeds
+device = get_device() # cuda > mps > cpu
+```
+
+
+## Used by
+
+- [persona-vectors](https://github.com/implicit-personalization/persona-vectors) — activation extraction and steering
+- [cues_attribution](https://github.com/implicit-personalization/io-analysis) — section-level ablation attribution
+- [persona-2-lora](https://github.com/implicit-personalization/persona-2-lora) — LoRA-based persona internalization

persona_data-0.1.0/README.md

@@ -0,0 +1,116 @@
+# persona-data
+
+[![Docs](https://img.shields.io/badge/docs-view-purple?logo=materialformkdocs)](https://implicit-personalization.github.io/persona-data/)
+
+Shared dataset loading, prompt formatting, and environment utilities for the [implicit-personalization](https://github.com/implicit-personalization) projects.
+
+## Overview
+
+`persona-data` provides the common dataset and prompt helpers used across the persona projects:
+
+- `SynthPersonaDataset` for persona profiles plus QA pairs
+- `PersonaGuessDataset` for turn-based persona games
+- prompt helpers for roleplay and multiple-choice evaluation
+- environment helpers for seeds, devices, and artifact paths
+
+## Installation
+
+Add as a uv git source in your project's `pyproject.toml`:
+
+```toml
+[project]
+dependencies = ["persona-data"]
+
+[tool.uv.sources]
+persona-data = { git = "ssh://git@github.com/implicit-personalization/persona-data.git" }
+```
+
+Then run `uv sync`.
+
+For local development alongside other repos, use an editable path source:
+
+```toml
+[tool.uv.sources]
+persona-data = { path = "../persona-data", editable = true }
+```
+
+## Package layout
+
+```
+src/persona_data/
+├── __init__.py
+├── synth_persona.py       # SynthPersonaDataset, PersonaDataset, PersonaData, QAPair, BiographySection
+├── persona_guess.py       # PersonaGuessDataset, GameRecord, Turn
+├── prompts.py             # format_roleplay_prompt, format_mc_question, format_messages
+└── environment.py         # load_env, set_seed, get_device, get_artifacts_dir
+```
+
+## Datasets
+
+Each dataset is a module with its own types and a loader that downloads from Hugging Face, cached via `HF_HOME`.
+
+### SynthPersona
+
+```python
+from persona_data.synth_persona import SynthPersonaDataset
+
+dataset = SynthPersonaDataset()
+
+persona = dataset[0]
+persona.name              # "Ethan Robinson"
+persona.templated_view    # short attribute-based system prompt
+persona.biography_view    # full biography text
+persona.sections          # list of BiographySection
+
+qa_pairs = dataset.get_qa(persona.id, type="implicit", difficulty=[1, 2])
+questions = dataset.questions(persona.id, type="explicit")
+```
+
+### PersonaGuess
+
+```python
+from persona_data.persona_guess import PersonaGuessDataset
+
+games = PersonaGuessDataset()
+game = games[0]
+turns = games.get_qa(game.game_id, player="A")
+questions = games.questions(game.game_id, player="B")
+```
+
+## Prompt formatting
+
+```python
+from persona_data.prompts import format_messages, format_roleplay_prompt
+
+system_prompt = format_roleplay_prompt(persona.biography_view)
+
+messages = [
+    {"role": "system", "content": system_prompt},
+    {"role": "user", "content": "Where did you grow up?"},
+    {"role": "assistant", "content": "I grew up in Little Rock, Arkansas."},
+]
+full_prompt, response_start_idx = format_messages(messages, tokenizer)
+```
+
+`format_roleplay_prompt` supports `mode="roleplay"` (default), `mode="conversational"`, and `mode="mc"`.
+
+`format_messages` handles tokenizers that do not support the `"system"` role (for example Gemma 2) by merging system content into the first user message.
+
+For multiple-choice evaluation, use `format_mc_question(qa)` and `mc_correct_letter(qa)`.
+
+## Environment helpers
+
+```python
+from persona_data.environment import load_env, set_seed, get_device, get_artifacts_dir
+
+load_env()            # loads .env from cwd (searches parent dirs)
+set_seed(1337)        # sets random, numpy, and torch seeds
+device = get_device() # cuda > mps > cpu
+```
+
+
+## Used by
+
+- [persona-vectors](https://github.com/implicit-personalization/persona-vectors) — activation extraction and steering
+- [cues_attribution](https://github.com/implicit-personalization/io-analysis) — section-level ablation attribution
+- [persona-2-lora](https://github.com/implicit-personalization/persona-2-lora) — LoRA-based persona internalization

persona_data-0.1.0/pyproject.toml

@@ -0,0 +1,16 @@
+[project]
+name = "persona-data"
+version = "0.1.0"
+description = "Shared dataset loading and prompt formatting for implicit-personalization projects"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "huggingface-hub>=0.30.0",
+    "python-dotenv>=1.0.0",
+    "torch>=2.0.0",
+    "numpy>=1.24.0",
+]
+
+[build-system]
+requires = ["uv_build>=0.11.3,<0.12"]
+build-backend = "uv_build"

persona_data-0.1.0/src/persona_data/environment.py

@@ -0,0 +1,42 @@
+import os
+import random
+from pathlib import Path
+
+import numpy as np
+import torch
+from dotenv import load_dotenv
+
+
+def load_env() -> None:
+    """Load environment variables from .env file.
+
+    Searches the current working directory and parent directories.
+    """
+    load_dotenv()
+
+
+def get_artifacts_dir() -> Path:
+    """Get the root artifacts directory."""
+    return Path(os.environ.get("ARTIFACTS_DIR", "artifacts"))
+
+
+def set_seed(seed: int) -> None:
+    """Set random seed for reproducibility."""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+    if torch.backends.mps.is_available():
+        torch.mps.manual_seed(seed)
+
+
+def get_device() -> torch.device:
+    """Determine the best available device."""
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    elif torch.backends.mps.is_available():
+        return torch.device("mps")
+    else:
+        return torch.device("cpu")

persona_data-0.1.0/src/persona_data/persona_guess.py

@@ -0,0 +1,71 @@
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterator, Literal
+
+from huggingface_hub import hf_hub_download
+
+
+@dataclass
+class Turn:
+    round: int
+    asker: Literal["A", "B"]
+    question: str
+    answer: str
+
+    def __repr__(self):
+        return f"Turn(round={self.round}, asker={self.asker!r})"
+
+
+@dataclass
+class GameRecord:
+    game_id: str
+    persona_a_id: str
+    persona_b_id: str
+    turns: list[Turn]
+
+    def __repr__(self):
+        return f"GameRecord(game_id={self.game_id!r}, turns={len(self.turns)})"
+
+
+class PersonaGuessDataset:
+    """PersonaGuess game dataset loaded from HuggingFace."""
+
+    DEFAULT_REPO = "implicit-personalization/persona-guess"
+
+    def __init__(self, hf_repo: str = DEFAULT_REPO) -> None:
+        path = Path(hf_hub_download(hf_repo, "games.jsonl", repo_type="dataset"))
+        with open(path) as f:
+            self._games: list[GameRecord] = [
+                GameRecord(
+                    game_id=d["game_id"],
+                    persona_a_id=d["persona_a_id"],
+                    persona_b_id=d["persona_b_id"],
+                    turns=[Turn(**t) for t in d["turns"]],
+                )
+                for d in (json.loads(line) for line in f)
+            ]
+        self._games_by_id: dict[str, GameRecord] = {g.game_id: g for g in self._games}
+
+    def __repr__(self) -> str:
+        return f"PersonaGuessDataset(n_games={len(self._games)})"
+
+    def __len__(self) -> int:
+        return len(self._games)
+
+    def __iter__(self) -> Iterator[GameRecord]:
+        return iter(self._games)
+
+    def __getitem__(self, idx: int) -> GameRecord:
+        return self._games[idx]
+
+    def get_qa(
+        self, game_id: str, player: Literal["A", "B"] | None = None
+    ) -> list[Turn]:
+        game = self._games_by_id[game_id]
+        return [t for t in game.turns if player is None or t.asker == player]
+
+    def questions(
+        self, game_id: str, player: Literal["A", "B"] | None = None
+    ) -> list[str]:
+        return [t.question for t in self.get_qa(game_id, player)]

persona_data-0.1.0/src/persona_data/prompts.py

@@ -0,0 +1,127 @@
+from typing import Literal
+
+from persona_data.synth_persona import QAPair
+
+_LETTERS = "ABCDE"
+
+BASE_ROLEPLAY_PROMPT = """\
+You are roleplaying as a specific person in a conversation.
+Stay fully in character. Be truthful to the profile below.
+Do not mention that you are an AI model.
+
+### Person profile:
+
+{persona}"""
+
+_CONVERSATIONAL_SUFFIX = "\n\nAnswer naturally and conversationally as this person."
+
+EMPTY_PERSONA_PLACEHOLDER = "Assistant"
+MC_ANSWER_ONLY_INSTRUCTION = "Answer only with the correct choice label (A, B, C, ...)."
+PromptMode = Literal["roleplay", "conversational", "mc"]
+
+
+def format_roleplay_prompt(
+    persona: str = EMPTY_PERSONA_PLACEHOLDER, mode: PromptMode = "roleplay"
+) -> str:
+    """System prompt using any persona view.
+
+    Args:
+        persona: The persona text (templated or biography view).
+        mode: Prompt style selector.
+            - "roleplay": plain persona prompt
+            - "conversational": persona prompt with a natural chat instruction
+            - "mc": persona prompt with a multiple-choice answer constraint
+    """
+    base = BASE_ROLEPLAY_PROMPT.format(persona=persona)
+    if mode == "conversational":
+        return base + _CONVERSATIONAL_SUFFIX
+    if mode == "mc":
+        return f"{base.rstrip()}\n\n{MC_ANSWER_ONLY_INSTRUCTION}"
+    return base
+
+
+def _format_mc_question_prompt(qa: QAPair) -> str:
+    lines = [qa.question, ""]
+    for i, choice in enumerate(qa.choices):
+        lines.append(f"{_LETTERS[i]}. {choice}")
+    return "\n".join(lines)
+
+
+def format_mc_question(qa: QAPair) -> str:
+    """Format an MC question with lettered choices for model evaluation."""
+    return f"{_format_mc_question_prompt(qa).rstrip()}\n\n{MC_ANSWER_ONLY_INSTRUCTION}"
+
+
+def mc_correct_letter(qa: QAPair) -> str:
+    """Return the letter (A–E) for the correct choice of an MC question."""
+    if qa.correct_choice_index is None:
+        raise ValueError(f"QAPair {qa.qid!r} has no correct_choice_index")
+    return _LETTERS[qa.correct_choice_index]
+
+
+def _supports_system_role(tokenizer) -> bool:
+    """Check if tokenizer's chat template supports the 'system' role."""
+    try:
+        tokenizer.apply_chat_template(
+            [{"role": "system", "content": "test"}],
+            tokenize=False,
+        )
+        return True
+    except Exception:
+        return False
+
+
+def normalize_messages(messages: list[dict[str, str]]) -> list[dict[str, str]]:
+    """Merge any leading system message into the first user message.
+
+    Only needed when the tokenizer's chat template doesn't support
+    the "system" role (e.g., Gemma 2).
+    """
+    if not messages or messages[0]["role"] != "system":
+        return messages
+    system_content = messages[0]["content"]
+    rest = list(messages[1:])
+    if rest and rest[0]["role"] == "user" and system_content:
+        rest[0] = {
+            "role": "user",
+            "content": f"{system_content}\n\n{rest[0]['content']}",
+        }
+    return rest
+
+
+def format_messages(messages: list[dict[str, str]], tokenizer) -> tuple[str, int]:
+    """Format a conversation for the model using its chat template.
+
+    Args:
+        messages: List of message dicts with "role" and "content" keys.
+                 Can include "system", "user", and "assistant" roles.
+        tokenizer: The tokenizer with chat template support.
+
+    Returns:
+        full_prompt: The full formatted prompt as a string.
+        response_start_idx: The token index of the first token in the last
+                            assistant message.
+    """
+    supports_system = _supports_system_role(tokenizer)
+    if not supports_system:
+        messages = normalize_messages(messages)
+
+    full_prompt = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=False
+    )
+
+    if len(messages) <= 1:
+        prompt_without_response = tokenizer.apply_chat_template(
+            messages, tokenize=True, add_generation_prompt=True, return_tensors="pt"
+        )
+    else:
+        prompt_without_response = tokenizer.apply_chat_template(
+            messages[:-1],
+            tokenize=True,
+            add_generation_prompt=True,
+            return_tensors="pt",
+        )
+
+    response_start_idx = prompt_without_response["input_ids"].shape[1]
+
+    return full_prompt, response_start_idx

persona_data-0.1.0/src/persona_data/synth_persona.py

@@ -0,0 +1,211 @@
+import json
+from collections import defaultdict
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Iterator, Literal
+
+# NOTE: The loader intentionally drops provenance/curation fields that are not used
+# by persona-data, persona-ui, or persona-vectors. We keep `evidence_sids` because
+# the ablation tooling still relies on it.
+# Dropped fields: design_notes, family_name, source_candidate_*, bank_id, axis,
+# curation_note, validation, evidence_claims, evidence_quotes.
+
+
+@dataclass
+class QAPair:
+    qid: str
+    type: Literal["explicit", "implicit"]
+    question: str
+    answer: str
+    difficulty: int  # 1 = easy, 2 = medium, 3 = hard
+    answer_format: str = ""  # "free_text" or "choice"
+    choices: list[str] = field(default_factory=list)
+    correct_choice_index: int | None = None
+    evidence_sids: list[str] = field(default_factory=list)
+    tags: list[str] = field(default_factory=list)
+
+    def __repr__(self):
+        return f"QAPair(qid={self.qid!r}, type={self.type!r}, difficulty={self.difficulty})"
+
+
+@dataclass
+class Statement:
+    sid: str
+    category: str
+    claim: str
+    support: list[dict] = field(default_factory=list)
+    confidence: str = ""
+
+    def __repr__(self):
+        return f"Statement(sid={self.sid!r}, category={self.category!r})"
+
+
+@dataclass
+class BiographySection:
+    """A semantically coherent segment of a persona biography."""
+
+    section_id: str
+    category: str  # e.g. "upbringing", "education", "career"
+    title: str
+    text: str  # Concatenated paragraph texts
+
+    def __repr__(self):
+        return f"BiographySection(id={self.section_id!r}, category={self.category!r})"
+
+
+@dataclass
+class PersonaData:
+    id: str
+    persona: dict
+    templated_view: str
+    biography_view: str
+    statements_view: str = ""
+    # transcript: str = "" # NOTE: This is not needed for our usecase for now
+    sections: list[BiographySection] = field(default_factory=list)
+    statements: list[Statement] = field(default_factory=list)
+
+    @property
+    def name(self) -> str:
+        return f"{self.persona['first_name']} {self.persona['last_name']}"
+
+    @property
+    def sections_by_id(self) -> dict[str, BiographySection]:
+        return {section.section_id: section for section in self.sections}
+
+    @property
+    def sections_by_category(self) -> dict[str, list[BiographySection]]:
+        grouped: dict[str, list[BiographySection]] = defaultdict(list)
+        for section in self.sections:
+            grouped[section.category].append(section)
+        return dict(grouped)
+
+    @property
+    def section_categories(self) -> list[str]:
+        """Unique section categories in order."""
+        return list(self.sections_by_category)
+
+    def get_section(self, section_id: str) -> BiographySection | None:
+        return self.sections_by_id.get(section_id)
+
+    def get_sections_by_category(self, category: str) -> list[BiographySection]:
+        return self.sections_by_category.get(category, [])
+
+    def __repr__(self):
+        return f"PersonaData(id={self.id!r}, name={self.name!r})"
+
+
+class PersonaDataset:
+    """Persona dataset loaded from local JSONL files."""
+
+    def __init__(self, personas_path: Path | str, qa_path: Path | str) -> None:
+        self._personas: list[PersonaData] = []
+        self._personas_by_id: dict[str, PersonaData] = {}
+        with open(personas_path) as f:
+            for line in f:
+                if not line.strip():
+                    continue
+                d = json.loads(line)
+                sections = [
+                    BiographySection(
+                        section_id=sec["section_id"],
+                        category=sec["category"],
+                        title=sec["title"],
+                        text="\n\n".join(p["text"] for p in sec.get("paragraphs", [])),
+                    )
+                    for sec in d.get("sections", [])
+                ]
+                persona = PersonaData(
+                    id=d["id"],
+                    persona=d["persona"],
+                    templated_view=d["templated_view"],
+                    biography_view=d.get("biography_view", ""),
+                    statements_view=d.get("statements_view", ""),
+                    sections=sections,
+                    statements=[
+                        Statement(
+                            sid=s["sid"],
+                            category=s["category"],
+                            claim=s["claim"],
+                            support=s.get("support", []),
+                            confidence=s.get("confidence", ""),
+                        )
+                        for s in d.get("statements", [])
+                    ],
+                )
+                self._personas.append(persona)
+                self._personas_by_id[persona.id] = persona
+
+        self._qa: dict[str, list[QAPair]] = defaultdict(list)
+        with open(qa_path) as f:
+            for line in f:
+                if not line.strip():
+                    continue
+                d = json.loads(line)
+                self._qa[d["id"]].append(
+                    QAPair(
+                        qid=d["qid"],
+                        type=d["type"],
+                        question=d["question"],
+                        answer=d["answer"],
+                        difficulty=d["difficulty"],
+                        answer_format=d.get("answer_format", "free_text"),
+                        choices=d.get("choices", []),
+                        correct_choice_index=d.get("correct_choice_index"),
+                        evidence_sids=d.get("evidence_sids", []),
+                        tags=d.get("tags", []),
+                    )
+                )
+
+    def __repr__(self) -> str:
+        return f"{type(self).__name__}(n_personas={len(self._personas)})"
+
+    def __len__(self) -> int:
+        return len(self._personas)
+
+    def __iter__(self) -> Iterator[PersonaData]:
+        return iter(self._personas)
+
+    def __getitem__(self, idx: int) -> PersonaData:
+        return self._personas[idx]
+
+    def get_persona(self, persona_id: str) -> PersonaData | None:
+        return self._personas_by_id.get(persona_id)
+
+    def get_qa(
+        self,
+        persona_id: str,
+        type: Literal["explicit", "implicit"] | None = None,
+        difficulty: int | list[int] | None = None,
+    ) -> list[QAPair]:
+        """Return QA pairs for a persona, optionally filtered by type and/or difficulty."""
+        pairs = self._qa.get(persona_id, [])
+        if type is not None:
+            pairs = [p for p in pairs if p.type == type]
+        if difficulty is not None:
+            levels = {difficulty} if isinstance(difficulty, int) else set(difficulty)
+            pairs = [p for p in pairs if p.difficulty in levels]
+        return pairs
+
+    def questions(
+        self,
+        persona_id: str,
+        type: Literal["explicit", "implicit"] | None = None,
+        difficulty: int | list[int] | None = None,
+    ) -> list[str]:
+        """Like get_qa but returns question strings only."""
+        return [qa.question for qa in self.get_qa(persona_id, type, difficulty)]
+
+
+class SynthPersonaDataset(PersonaDataset):
+    """SynthPersona dataset loaded from HuggingFace."""
+
+    def __init__(self, hf_repo: str = "implicit-personalization/synth-persona") -> None:
+        from huggingface_hub import hf_hub_download
+
+        # HF Hub caches locally under HF_HOME so repeat runs are instant.
+        super().__init__(
+            personas_path=hf_hub_download(
+                hf_repo, "dataset_personas.jsonl", repo_type="dataset"
+            ),
+            qa_path=hf_hub_download(hf_repo, "dataset_qa.jsonl", repo_type="dataset"),
+        )