Cleaned up code abastracted away in persona-vector

app.py

@@ -1,14 +1,11 @@
 import os
-from pathlib import Path
 
 import streamlit as st
 from dotenv import load_dotenv
 
-# Load .env early so DEFAULT_MODEL / REMOTE_DEFAULT_MODEL can be overridden via env
-load_dotenv(Path(__file__).parent / ".env")
-
 from utils.helpers import DATASET_SOURCES
 
+load_dotenv()
 DEFAULT_MODEL = os.environ.get("DEFAULT_MODEL", "google/gemma-2-2b-it")
 REMOTE_DEFAULT_MODEL = os.environ.get("REMOTE_DEFAULT_MODEL", "google/gemma-2-9b-it")
 

tabs/extract.py

@@ -1,7 +1,7 @@
 import streamlit as st
+from persona_vectors.extraction import run_extraction
 
 from utils.datasets import load_dataset
-from utils.extraction import run_extraction
 from utils.helpers import (
     PROMPT_VARIANTS,
     persona_label,
@@ -151,6 +151,16 @@ def render_extract_tab(remote: bool, model_name: str, dataset_source: str) -> No
     status_box = st.empty()
     status_box.info("Extraction in progress...")
     progress = st.progress(0, text="Preparing extraction...")
+    ndif_status_box = st.empty()  # shows live NDIF job status when remote=True
+
+    _STATUS_ICONS = {
+        "RECEIVED": "◉", "QUEUED": "◎", "DISPATCHED": "◈",
+        "RUNNING": "●", "COMPLETED": "✓", "ERROR": "✗",
+    }
+
+    def _on_ndif_status(job_id: str, status_name: str, description: str) -> None:
+        icon = _STATUS_ICONS.get(status_name, "•")
+        ndif_status_box.caption(f"{icon} `{job_id}` **{status_name}** — {description}")
 
     with st.spinner("Loading model..."):
         model = cached_model(model_name=model_name, remote=remote)
@@ -174,6 +184,7 @@ def render_extract_tab(remote: bool, model_name: str, dataset_source: str) -> No
                     qa_pairs=qa_pairs,
                     variants=[variant],
                     remote=remote,
+                    on_status=_on_ndif_status if remote else None,
                 )
                 results.extend(variant_results)
                 step += 1
@@ -184,6 +195,7 @@ def render_extract_tab(remote: bool, model_name: str, dataset_source: str) -> No
         return
     finally:
         progress.empty()
+        ndif_status_box.empty()
 
     status_box.success("Extraction complete")
     st.success(f"Saved {len(results)} artifact set(s)")
@@ -191,5 +203,5 @@ def render_extract_tab(remote: bool, model_name: str, dataset_source: str) -> No
     for result in results:
         st.markdown(
             f"- **{result.persona_name}** · {prompt_variant_label(result.variant)}: "
-            f"{result.n_questions} questions, {result.n_layers} layers, {result.d_model} hidden size"
+            f"{result.n_questions} questions"
         )

utils/artifacts.py

@@ -7,17 +7,12 @@ import torch
 from persona_vectors.activation_io import (
     load_activation_metadata,
     load_per_question_vectors,
+    model_dir_name,
 )
 
 logger = logging.getLogger(__name__)
 
 
-def model_dir_name(model_name: str) -> str:
-    """Encode a model name for use in artifact paths."""
-
-    return model_name.replace("/", "__")
-
-
 def list_available_personas(
     artifacts_root: str | Path,
     model_name: str,

utils/chat.py

@@ -8,13 +8,13 @@ from nnterp import StandardizedTransformer
 
 logger = logging.getLogger(__name__)
 
-from persona_data.synth_persona import PersonaData
 from persona_data.prompts import (
-    format_empty_persona_prompt,
     format_biography_prompt,
+    format_empty_persona_prompt,
     format_templated_prompt,
     normalize_messages,
 )
+from persona_data.synth_persona import PersonaData
 
 SystemPromptMode = Literal["empty", "templated", "biography", "custom"]
 

utils/chat_export.py

@@ -3,8 +3,8 @@ from datetime import datetime, timezone
 from pathlib import Path
 
 from persona_data.environment import get_artifacts_dir
+from persona_vectors.activation_io import model_dir_name
 
-from utils.artifacts import model_dir_name
 from utils.helpers import slugify
 
 

utils/datasets.py

@@ -5,10 +5,10 @@ from tempfile import mkdtemp
 from typing import Any
 
 import streamlit as st
+from persona_data.synth_persona import PersonaDataset as LocalPersonaDataset
 from persona_data.synth_persona import SynthPersonaDataset
 
 from .helpers import DATASET_SOURCES
-from .local_dataset import LocalPersonaDataset
 
 
 @st.cache_resource(show_spinner=False)

utils/extraction.py

@@ -1,151 +0,0 @@
-import gc
-import logging
-from dataclasses import dataclass
-
-import torch
-from nnterp import StandardizedTransformer
-
-logger = logging.getLogger(__name__)
-
-from persona_data.environment import get_artifacts_dir
-from persona_data.synth_persona import PersonaData, QAPair
-from persona_vectors.activation_io import save_per_question_vectors
-from persona_vectors.activations import extract_activations
-from persona_data.prompts import (
-    format_biography_prompt,
-    format_messages,
-    format_templated_prompt,
-)
-
-
-@dataclass
-class VariantExtractionResult:
-    variant: str
-    output_dir: str
-    n_questions: int
-    n_layers: int
-    d_model: int
-    persona_name: str = ""
-
-
-def _prepare_inputs(
-    tokenizer: object,
-    system_prompt: str,
-    qa_pairs: list[QAPair],
-) -> tuple[list[str], list[torch.Tensor], list[str]]:
-    """Format QA pairs into tokenized prompts with answer-token masks.
-
-    Args:
-        tokenizer: HuggingFace-compatible tokenizer from the model.
-        system_prompt: System prompt to prepend to each conversation.
-        qa_pairs: List of question-answer pairs to format.
-
-    Returns:
-        A tuple of (full_texts, token_masks, questions) where full_texts are
-        the rendered prompt strings, token_masks are boolean tensors marking
-        answer tokens, and questions are the raw question strings.
-    """
-    full_texts: list[str] = []
-    token_masks: list[torch.Tensor] = []
-    questions: list[str] = []
-
-    for qa in qa_pairs:
-        messages = [
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": qa.question},
-            {"role": "assistant", "content": qa.answer},
-        ]
-        full_prompt, answer_start = format_messages(messages, tokenizer)
-        seq_len = tokenizer(full_prompt, return_tensors="pt").input_ids.shape[1]
-
-        full_texts.append(full_prompt)
-        token_masks.append(torch.arange(seq_len) >= answer_start)
-        questions.append(qa.question)
-
-    return full_texts, token_masks, questions
-
-
-def run_extraction(
-    model: StandardizedTransformer,
-    model_name: str,
-    persona: PersonaData,
-    qa_pairs: list[QAPair],
-    variants: list[str],
-    remote: bool,
-) -> list[VariantExtractionResult]:
-    """Run activation extraction and save outputs for selected variants.
-
-    Args:
-        model: Loaded standardized nnterp model.
-        model_name: HuggingFace model identifier used for artifact paths.
-        persona: The persona whose QA pairs are being extracted.
-        qa_pairs: Question-answer pairs to run extraction on.
-        variants: Prompt variants to extract (e.g. ``"templated"``, ``"biography"``).
-        remote: Whether to execute on NDIF.
-
-    Returns:
-        A list of extraction results, one per variant.
-
-    Raises:
-        ValueError: If ``qa_pairs`` is empty or an unsupported variant is given.
-    """
-    if not qa_pairs:
-        raise ValueError("No QA pairs selected for extraction")
-
-    tokenizer = model.tokenizer
-    activations_dir = get_artifacts_dir() / "activations"
-
-    system_prompt_by_variant = {
-        "templated": format_templated_prompt(persona.templated_prompt),
-        "biography": format_biography_prompt(persona.biography_md),
-    }
-
-    results: list[VariantExtractionResult] = []
-
-    for variant in variants:
-        if variant not in system_prompt_by_variant:
-            raise ValueError(f"Unsupported variant: {variant}")
-
-        full_texts, token_masks, questions = _prepare_inputs(
-            tokenizer=tokenizer,
-            system_prompt=system_prompt_by_variant[variant],
-            qa_pairs=qa_pairs,
-        )
-
-        per_question_vectors = extract_activations(
-            model=model,
-            full_texts=full_texts,
-            token_masks=token_masks,
-            remote=remote,
-        )
-
-        artifact_dir = save_per_question_vectors(
-            root_dir=activations_dir,
-            model_name=model_name,
-            prompt_variant=variant,
-            persona_id=persona.id,
-            persona_name=persona.name,
-            per_question_vectors=per_question_vectors,
-            questions=questions,
-        )
-
-        results.append(
-            VariantExtractionResult(
-                variant=variant,
-                output_dir=str(artifact_dir),
-                n_questions=per_question_vectors.shape[0],
-                n_layers=per_question_vectors.shape[1],
-                d_model=per_question_vectors.shape[2],
-                persona_name=persona.name,
-            )
-        )
-
-        # Free activation tensors between variants to keep memory bounded.
-        del per_question_vectors, full_texts, token_masks
-        gc.collect()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        if hasattr(torch, "mps") and hasattr(torch.mps, "empty_cache"):
-            torch.mps.empty_cache()
-
-    return results

utils/helpers.py

@@ -1,4 +1,5 @@
 from persona_data.synth_persona import PersonaData
+from persona_vectors.extraction import SUPPORTED_VARIANTS
 
 # Variant key -> human-readable label mapping
 VARIANT_LABELS = {
@@ -9,7 +10,7 @@ VARIANT_LABELS = {
 }
 
 # Variants that correspond to actual system prompts (excludes "empty")
-PROMPT_VARIANTS = ["templated", "biography"]
+PROMPT_VARIANTS = list(SUPPORTED_VARIANTS)
 
 # For selectbox options: list of labels in definition order
 MODE_LABELS = list(VARIANT_LABELS.values())

utils/local_dataset.py

@@ -1,72 +0,0 @@
-import json
-from collections import defaultdict
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Iterator, Literal
-
-from persona_data.synth_persona import PersonaData, QAPair
-
-
-@dataclass
-class LocalPersonaDataset:
-    """Dataset loaded from local JSONL files."""
-
-    personas_path: Path
-    qa_path: Path
-
-    def __post_init__(self) -> None:
-        with self.personas_path.open() as f:
-            self._personas: list[PersonaData] = []
-            for line in f:
-                if not line.strip():
-                    continue
-                data = json.loads(line)
-                self._personas.append(
-                    PersonaData(
-                        id=data["id"],
-                        persona=data["persona"],
-                        templated_prompt=data["templated_prompt"],
-                        biography_md=data["biography_md"],
-                    )
-                )
-
-        self._qa: dict[str, list[QAPair]] = defaultdict(list)
-        with self.qa_path.open() as f:
-            for line in f:
-                if not line.strip():
-                    continue
-                data = json.loads(line)
-                self._qa[data["id"]].append(
-                    QAPair(
-                        qid=data["qid"],
-                        type=data["type"],
-                        question=data["question"],
-                        answer=data["answer"],
-                        difficulty=data["difficulty"],
-                    )
-                )
-
-    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_qa(
-        self,
-        persona_id: str,
-        type: Literal["explicit", "implicit"] | None = None,
-        difficulty: int | list[int] | None = None,
-    ) -> list[QAPair]:
-        pairs = self._qa.get(persona_id, [])
-        if type is not None:
-            pairs = [pair for pair in pairs if pair.type == type]
-
-        if difficulty is not None:
-            levels = {difficulty} if isinstance(difficulty, int) else set(difficulty)
-            pairs = [pair for pair in pairs if pair.difficulty in levels]
-
-        return pairs