Push

Dockerfile

@@ -13,7 +13,7 @@ COPY requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
 
 # Copy the data to /data, the HF persistent storage. We do this after pip install to avoid
-# re-installing dependencies if the data changes.
+# re-installing dependencies if the data changes, which is likely more often.
 WORKDIR /
 COPY /data /data
 WORKDIR /server
@@ -27,4 +27,7 @@ COPY /web/blueprint/build ./web/blueprint/build
 # Copy python files.
 COPY /src ./src/
 
+# Copy the entrypoint file.
+COPY docker_entrypoint.sh .
+
 CMD ["uvicorn", "src.server:app", "--host", "0.0.0.0", "--port", "5432"]

requirements.txt

@@ -1,6 +1,6 @@
 aiohttp==3.8.4 ; python_version >= "3.9" and python_version < "3.10"
 aiosignal==1.3.1 ; python_version >= "3.9" and python_version < "3.10"
-anyio==3.7.0 ; python_version >= "3.9" and python_version < "3.10"
+anyio==3.7.1 ; python_version >= "3.9" and python_version < "3.10"
 async-timeout==4.0.2 ; python_version >= "3.9" and python_version < "3.10"
 attrs==23.1.0 ; python_version >= "3.9" and python_version < "3.10"
 blis==0.7.9 ; python_version >= "3.9" and python_version < "3.10"
@@ -12,43 +12,43 @@ click==8.1.3 ; python_version >= "3.9" and python_version < "3.10"
 cloudpickle==2.2.1 ; python_version >= "3.9" and python_version < "3.10"
 cohere==3.10.0 ; python_version >= "3.9" and python_version < "3.10"
 colorama==0.4.6 ; python_version >= "3.9" and python_version < "3.10" and (platform_system == "Windows" or sys_platform == "win32")
-confection==0.0.4 ; python_version >= "3.9" and python_version < "3.10"
+confection==0.1.0 ; python_version >= "3.9" and python_version < "3.10"
 cymem==2.0.7 ; python_version >= "3.9" and python_version < "3.10"
 cytoolz==0.12.1 ; python_version >= "3.9" and python_version < "3.10"
-dask==2023.5.1 ; python_version >= "3.9" and python_version < "3.10"
-datasets==2.12.0 ; python_version >= "3.9" and python_version < "3.10"
+dask==2023.6.1 ; python_version >= "3.9" and python_version < "3.10"
+datasets==2.13.1 ; python_version >= "3.9" and python_version < "3.10"
 decorator==5.1.1 ; python_version >= "3.9" and python_version < "3.10"
 dill==0.3.6 ; python_version >= "3.9" and python_version < "3.10"
-distributed==2023.5.1 ; python_version >= "3.9" and python_version < "3.10"
-duckdb==0.8.0 ; python_version >= "3.9" and python_version < "3.10"
+distributed==2023.6.1 ; python_version >= "3.9" and python_version < "3.10"
+duckdb==0.8.1 ; python_version >= "3.9" and python_version < "3.10"
 email-reply-parser==0.5.12 ; python_version >= "3.9" and python_version < "3.10"
-exceptiongroup==1.1.1 ; python_version >= "3.9" and python_version < "3.10"
-fastapi==0.95.2 ; python_version >= "3.9" and python_version < "3.10"
-filelock==3.12.0 ; python_version >= "3.9" and python_version < "3.10"
+exceptiongroup==1.1.2 ; python_version >= "3.9" and python_version < "3.10"
+fastapi==0.98.0 ; python_version >= "3.9" and python_version < "3.10"
+filelock==3.12.2 ; python_version >= "3.9" and python_version < "3.10"
 floret==0.10.3 ; python_version >= "3.9" and python_version < "3.10"
 frozenlist==1.3.3 ; python_version >= "3.9" and python_version < "3.10"
-fsspec==2023.5.0 ; python_version >= "3.9" and python_version < "3.10"
-fsspec[http]==2023.5.0 ; python_version >= "3.9" and python_version < "3.10"
-gcsfs==2023.5.0 ; python_version >= "3.9" and python_version < "3.10"
-google-api-core==2.11.0 ; python_version >= "3.9" and python_version < "3.10"
-google-api-python-client==2.88.0 ; python_version >= "3.9" and python_version < "3.10"
+fsspec==2023.6.0 ; python_version >= "3.9" and python_version < "3.10"
+fsspec[http]==2023.6.0 ; python_version >= "3.9" and python_version < "3.10"
+gcsfs==2023.6.0 ; python_version >= "3.9" and python_version < "3.10"
+google-api-core==2.11.1 ; python_version >= "3.9" and python_version < "3.10"
+google-api-python-client==2.92.0 ; python_version >= "3.9" and python_version < "3.10"
 google-auth-httplib2==0.1.0 ; python_version >= "3.9" and python_version < "3.10"
 google-auth-oauthlib==1.0.0 ; python_version >= "3.9" and python_version < "3.10"
-google-auth==2.19.1 ; python_version >= "3.9" and python_version < "3.10"
-google-cloud-core==2.3.2 ; python_version >= "3.9" and python_version < "3.10"
-google-cloud-storage==2.9.0 ; python_version >= "3.9" and python_version < "3.10"
+google-auth==2.21.0 ; python_version >= "3.9" and python_version < "3.10"
+google-cloud-core==2.3.3 ; python_version >= "3.9" and python_version < "3.10"
+google-cloud-storage==2.10.0 ; python_version >= "3.9" and python_version < "3.10"
 google-crc32c==1.5.0 ; python_version >= "3.9" and python_version < "3.10"
 google-resumable-media==2.5.0 ; python_version >= "3.9" and python_version < "3.10"
-googleapis-common-protos==1.59.0 ; python_version >= "3.9" and python_version < "3.10"
+googleapis-common-protos==1.59.1 ; python_version >= "3.9" and python_version < "3.10"
 h11==0.14.0 ; python_version >= "3.9" and python_version < "3.10"
 httplib2==0.22.0 ; python_version >= "3.9" and python_version < "3.10"
 httptools==0.5.0 ; python_version >= "3.9" and python_version < "3.10"
 huggingface-hub==0.15.1 ; python_version >= "3.9" and python_version < "3.10"
 idna==3.4 ; python_version >= "3.9" and python_version < "3.10"
-importlib-metadata==6.6.0 ; python_version >= "3.9" and python_version < "3.10"
-jellyfish==0.11.2 ; python_version >= "3.9" and python_version < "3.10"
+importlib-metadata==6.7.0 ; python_version >= "3.9" and python_version < "3.10"
+jellyfish==1.0.0 ; python_version >= "3.9" and python_version < "3.10"
 jinja2==3.1.2 ; python_version >= "3.9" and python_version < "3.10"
-joblib==1.2.0 ; python_version >= "3.9" and python_version < "3.10"
+joblib==1.3.1 ; python_version >= "3.9" and python_version < "3.10"
 langcodes==3.3.0 ; python_version >= "3.9" and python_version < "3.10"
 locket==1.0.0 ; python_version >= "3.9" and python_version < "3.10"
 markupsafe==2.1.3 ; python_version >= "3.9" and python_version < "3.10"
@@ -59,24 +59,24 @@ multiprocess==0.70.14 ; python_version >= "3.9" and python_version < "3.10"
 murmurhash==1.0.9 ; python_version >= "3.9" and python_version < "3.10"
 networkx==3.1 ; python_version >= "3.9" and python_version < "3.10"
 nltk==3.8.1 ; python_version >= "3.9" and python_version < "3.10"
-numpy==1.24.3 ; python_version >= "3.9" and python_version < "3.10"
+numpy==1.25.0 ; python_version >= "3.9" and python_version < "3.10"
 oauthlib==3.2.2 ; python_version >= "3.9" and python_version < "3.10"
 openai-function-call==0.0.5 ; python_version >= "3.9" and python_version < "3.10"
 openai==0.27.8 ; python_version >= "3.9" and python_version < "3.10"
-orjson==3.9.0 ; python_version >= "3.9" and python_version < "3.10"
+orjson==3.9.1 ; python_version >= "3.9" and python_version < "3.10"
 packaging==23.1 ; python_version >= "3.9" and python_version < "3.10"
-pandas==2.0.2 ; python_version >= "3.9" and python_version < "3.10"
+pandas==2.0.3 ; python_version >= "3.9" and python_version < "3.10"
 partd==1.4.0 ; python_version >= "3.9" and python_version < "3.10"
-pathy==0.10.1 ; python_version >= "3.9" and python_version < "3.10"
+pathy==0.10.2 ; python_version >= "3.9" and python_version < "3.10"
 pillow==9.5.0 ; python_version >= "3.9" and python_version < "3.10"
 preshed==3.0.8 ; python_version >= "3.9" and python_version < "3.10"
-protobuf==4.23.2 ; python_version >= "3.9" and python_version < "3.10"
+protobuf==4.23.3 ; python_version >= "3.9" and python_version < "3.10"
 psutil==5.9.5 ; python_version >= "3.9" and python_version < "3.10"
 pyarrow==9.0.0 ; python_version >= "3.9" and python_version < "3.10"
 pyasn1-modules==0.3.0 ; python_version >= "3.9" and python_version < "3.10"
 pyasn1==0.5.0 ; python_version >= "3.9" and python_version < "3.10"
-pydantic==1.10.9 ; python_version >= "3.9" and python_version < "3.10"
-pyparsing==3.0.9 ; python_version >= "3.9" and python_version < "3.10"
+pydantic==1.10.11 ; python_version >= "3.9" and python_version < "3.10"
+pyparsing==3.1.0 ; python_version >= "3.9" and python_version < "3.10"
 pyphen==0.14.0 ; python_version >= "3.9" and python_version < "3.10"
 python-dateutil==2.8.2 ; python_version >= "3.9" and python_version < "3.10"
 python-dotenv==1.0.0 ; python_version >= "3.9" and python_version < "3.10"
@@ -85,10 +85,10 @@ pyyaml==6.0 ; python_version >= "3.9" and python_version < "3.10"
 regex==2023.6.3 ; python_version >= "3.9" and python_version < "3.10"
 requests-oauthlib==1.3.1 ; python_version >= "3.9" and python_version < "3.10"
 requests==2.31.0 ; python_version >= "3.9" and python_version < "3.10"
-responses==0.18.0 ; python_version >= "3.9" and python_version < "3.10"
 rsa==4.9 ; python_version >= "3.9" and python_version < "3.10"
-scikit-learn==1.2.2 ; python_version >= "3.9" and python_version < "3.10"
-scipy==1.10.1 ; python_version >= "3.9" and python_version < "3.10"
+safetensors==0.3.1 ; python_version >= "3.9" and python_version < "3.10"
+scikit-learn==1.3.0 ; python_version >= "3.9" and python_version < "3.10"
+scipy==1.11.1 ; python_version >= "3.9" and python_version < "3.10"
 sentence-transformers==2.2.2 ; python_version >= "3.9" and python_version < "3.10"
 sentencepiece==0.1.99 ; python_version >= "3.9" and python_version < "3.10"
 setuptools==65.7.0 ; python_version >= "3.9" and python_version < "3.10"
@@ -98,11 +98,12 @@ sniffio==1.3.0 ; python_version >= "3.9" and python_version < "3.10"
 sortedcontainers==2.4.0 ; python_version >= "3.9" and python_version < "3.10"
 spacy-legacy==3.0.12 ; python_version >= "3.9" and python_version < "3.10"
 spacy-loggers==1.0.4 ; python_version >= "3.9" and python_version < "3.10"
-spacy==3.5.3 ; python_version >= "3.9" and python_version < "3.10"
+spacy==3.5.4 ; python_version >= "3.9" and python_version < "3.10"
 srsly==2.4.6 ; python_version >= "3.9" and python_version < "3.10"
 starlette==0.27.0 ; python_version >= "3.9" and python_version < "3.10"
 sympy==1.12 ; python_version >= "3.9" and python_version < "3.10"
-tblib==1.7.0 ; python_version >= "3.9" and python_version < "3.10"
+tblib==2.0.0 ; python_version >= "3.9" and python_version < "3.10"
+tenacity==8.2.2 ; python_version >= "3.9" and python_version < "3.10"
 textacy==0.13.0 ; python_version >= "3.9" and python_version < "3.10"
 thinc==8.1.10 ; python_version >= "3.9" and python_version < "3.10"
 threadpoolctl==3.1.0 ; python_version >= "3.9" and python_version < "3.10"
@@ -112,16 +113,16 @@ torch==2.0.1 ; python_version >= "3.9" and python_version < "3.10"
 torchvision==0.15.2 ; python_version >= "3.9" and python_version < "3.10"
 tornado==6.3.2 ; python_version >= "3.9" and python_version < "3.10"
 tqdm==4.65.0 ; python_version >= "3.9" and python_version < "3.10"
-transformers==4.29.2 ; python_version >= "3.9" and python_version < "3.10"
-typer==0.7.0 ; python_version >= "3.9" and python_version < "3.10"
-types-psutil==5.9.5.13 ; python_version >= "3.9" and python_version < "3.10"
-typing-extensions==4.6.3 ; python_version >= "3.9" and python_version < "3.10"
+transformers==4.30.2 ; python_version >= "3.9" and python_version < "3.10"
+typer==0.9.0 ; python_version >= "3.9" and python_version < "3.10"
+types-psutil==5.9.5.15 ; python_version >= "3.9" and python_version < "3.10"
+typing-extensions==4.7.1 ; python_version >= "3.9" and python_version < "3.10"
 tzdata==2023.3 ; python_version >= "3.9" and python_version < "3.10"
 uritemplate==4.1.1 ; python_version >= "3.9" and python_version < "3.10"
 urllib3==1.26.16 ; python_version >= "3.9" and python_version < "3.10"
-uvicorn[standard]==0.20.0 ; python_version >= "3.9" and python_version < "3.10"
+uvicorn[standard]==0.22.0 ; python_version >= "3.9" and python_version < "3.10"
 uvloop==0.17.0 ; (sys_platform != "win32" and sys_platform != "cygwin") and platform_python_implementation != "PyPy" and python_version >= "3.9" and python_version < "3.10"
-wasabi==1.1.1 ; python_version >= "3.9" and python_version < "3.10"
+wasabi==1.1.2 ; python_version >= "3.9" and python_version < "3.10"
 watchfiles==0.19.0 ; python_version >= "3.9" and python_version < "3.10"
 websockets==11.0.3 ; python_version >= "3.9" and python_version < "3.10"
 xxhash==3.2.0 ; python_version >= "3.9" and python_version < "3.10"

src/concepts/concept.py

@@ -1,23 +1,37 @@
 """Defines the concept and the concept models."""
+import dataclasses
 import random
-from typing import Iterable, Literal, Optional, Union
+from enum import Enum
+from typing import Callable, Iterable, Literal, Optional, Union
 
 import numpy as np
+from joblib import Parallel, delayed
 from pydantic import BaseModel, validator
+from scipy.interpolate import interp1d
+from sklearn.base import BaseEstimator, clone
 from sklearn.exceptions import NotFittedError
 from sklearn.linear_model import LogisticRegression
+from sklearn.metrics import precision_recall_curve, roc_auc_score
+from sklearn.model_selection import KFold
 
 from ..db_manager import get_dataset
 from ..embeddings.embedding import get_embed_fn
 from ..schema import Path, RichData, SignalInputType, normalize_path
-from ..signals.signal import TextEmbeddingSignal, get_signal_cls
+from ..signals.signal import EMBEDDING_KEY, TextEmbeddingSignal, get_signal_cls
 from ..utils import DebugTimer
 
 LOCAL_CONCEPT_NAMESPACE = 'local'
 
 # Number of randomly sampled negative examples to use for training. This is used to obtain a more
 # balanced model that works with a specific dataset.
-DEFAULT_NUM_NEG_EXAMPLES = 300
+DEFAULT_NUM_NEG_EXAMPLES = 100
+
+# The maximum number of cross-validation models to train.
+MAX_NUM_CROSS_VAL_MODELS = 15
+# The β weight to use for the F-beta score: https://scikit-learn.org/stable/modules/generated/sklearn.metrics.fbeta_score.html
+# β = 0.5 means we value precision 2x as much as recall.
+# β = 2 means we value recall 2x as much as precision.
+F_BETA_WEIGHT = 0.5
 
 
 class ConceptColumnInfo(BaseModel):
@@ -29,6 +43,13 @@ class ConceptColumnInfo(BaseModel):
   # Path holding the text to use for negative examples.
   path: Path
 
+  @validator('path')
+  def _path_points_to_text_field(cls, path: Path) -> Path:
+    if path[-1] == EMBEDDING_KEY:
+      raise ValueError(
+        f'The path should point to the text field, not its embedding field. Provided path: {path}')
+    return path
+
   num_negative_examples = DEFAULT_NUM_NEG_EXAMPLES
 
 
@@ -71,7 +92,7 @@ class Example(ExampleIn):
 class Concept(BaseModel):
   """A concept is a collection of examples."""
   # The namespace of the concept.
-  namespace: str = LOCAL_CONCEPT_NAMESPACE
+  namespace: str
   # The name of the concept.
   concept_name: str
   # The type of the data format that this concept represents.
@@ -79,6 +100,8 @@ class Concept(BaseModel):
   data: dict[str, Example]
   version: int = 0
 
+  description: Optional[str] = None
+
   def drafts(self) -> list[DraftId]:
     """Gets all the drafts for the concept."""
     drafts: set[DraftId] = set([DRAFT_MAIN])  # Always return the main draft.
@@ -88,39 +111,141 @@ class Concept(BaseModel):
     return list(sorted(drafts))
 
 
-class LogisticEmbeddingModel(BaseModel):
-  """A model that uses logistic regression with embeddings."""
+class OverallScore(str, Enum):
+  """Enum holding the overall score."""
+  NOT_GOOD = 'not_good'
+  OK = 'ok'
+  GOOD = 'good'
+  VERY_GOOD = 'very_good'
+  GREAT = 'great'
 
-  class Config:
-    arbitrary_types_allowed = True
-    underscore_attrs_are_private = True
 
-  version: int = -1
+def _get_overall_score(f1_score: float) -> OverallScore:
+  if f1_score < 0.5:
+    return OverallScore.NOT_GOOD
+  if f1_score < 0.8:
+    return OverallScore.OK
+  if f1_score < 0.9:
+    return OverallScore.GOOD
+  if f1_score < 0.95:
+    return OverallScore.VERY_GOOD
+  return OverallScore.GREAT
+
+
+class ConceptMetrics(BaseModel):
+  """Metrics for a concept."""
+  # The average F1 score for the concept computed using cross validation.
+  f1: float
+  precision: float
+  recall: float
+  roc_auc: float
+  overall: OverallScore
 
-  # The following fields are excluded from JSON serialization, but still pickleable.
-  # See `notebooks/Toxicity.ipynb` for an example of training a concept model.
-  _model: LogisticRegression = LogisticRegression(
-    class_weight=None, C=30, tol=1e-5, warm_start=True, max_iter=1_000, n_jobs=-1)
+
+@dataclasses.dataclass
+class LogisticEmbeddingModel:
+  """A model that uses logistic regression with embeddings."""
+
+  _metrics: Optional[ConceptMetrics] = None
+  _threshold: float = 0.5
+
+  def __post_init__(self) -> None:
+    # See `notebooks/Toxicity.ipynb` for an example of training a concept model.
+    self._model = LogisticRegression(
+      class_weight=None, C=30, tol=1e-5, warm_start=True, max_iter=1_000, n_jobs=-1)
 
   def score_embeddings(self, embeddings: np.ndarray) -> np.ndarray:
     """Get the scores for the provided embeddings."""
     try:
-      return self._model.predict_proba(embeddings)[:, 1]
+      y_probs = self._model.predict_proba(embeddings)[:, 1]
+      # Map [0, threshold, 1] to [0, 0.5, 1].
+      interpolate_fn = interp1d([0, self._threshold, 1], [0, 0.4999, 1])
+      return interpolate_fn(y_probs)
     except NotFittedError:
       return np.random.rand(len(embeddings))
 
-  def fit(self, embeddings: np.ndarray, labels: list[bool], sample_weights: list[float]) -> None:
+  def _setup_training(
+      self, X_train: np.ndarray, y_train: list[bool],
+      implicit_negatives: Optional[np.ndarray]) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    num_pos_labels = len([y for y in y_train if y])
+    num_neg_labels = len([y for y in y_train if not y])
+    sample_weights = [(1.0 / num_pos_labels if y else 1.0 / num_neg_labels) for y in y_train]
+
+    if implicit_negatives is not None:
+      num_implicit_labels = len(implicit_negatives)
+      implicit_labels = [False] * num_implicit_labels
+      X_train = np.concatenate([implicit_negatives, X_train])
+      y_train = np.concatenate([implicit_labels, y_train])
+      sample_weights = [1.0 / num_implicit_labels] * num_implicit_labels + sample_weights
+
+    # Normalize sample weights to sum to the number of training examples.
+    weights = np.array(sample_weights)
+    weights *= (X_train.shape[0] / np.sum(weights))
+    return X_train, np.array(y_train), weights
+
+  def fit(self, embeddings: np.ndarray, labels: list[bool],
+          implicit_negatives: Optional[np.ndarray]) -> None:
     """Fit the model to the provided embeddings and labels."""
-    if len(set(labels)) < 2:
+    label_set = set(labels)
+    if implicit_negatives is not None:
+      label_set.add(False)
+    if len(label_set) < 2:
       return
     if len(labels) != len(embeddings):
       raise ValueError(
         f'Length of embeddings ({len(embeddings)}) must match length of labels ({len(labels)})')
-    if len(sample_weights) != len(labels):
-      raise ValueError(
-        f'Length of sample_weights ({len(sample_weights)}) must match length of labels '
-        f'({len(labels)})')
-    self._model.fit(embeddings, labels, sample_weights)
+    X_train, y_train, sample_weights = self._setup_training(embeddings, labels, implicit_negatives)
+    self._model.fit(X_train, y_train, sample_weights)
+    self._metrics, self._threshold = self._compute_metrics(embeddings, labels, implicit_negatives)
+
+  def _compute_metrics(
+      self, embeddings: np.ndarray, labels: list[bool],
+      implicit_negatives: Optional[np.ndarray]) -> tuple[Optional[ConceptMetrics], float]:
+    """Return the concept metrics."""
+    labels = np.array(labels)
+    n_splits = min(len(labels), MAX_NUM_CROSS_VAL_MODELS)
+    fold = KFold(n_splits, shuffle=True, random_state=42)
+
+    def _fit_and_score(model: BaseEstimator, X_train: np.ndarray, y_train: np.ndarray,
+                       sample_weights: np.ndarray, X_test: np.ndarray,
+                       y_test: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+      if len(set(y_train)) < 2:
+        return np.array([]), np.array([])
+      model.fit(X_train, y_train, sample_weights)
+      y_pred = model.predict_proba(X_test)[:, 1]
+      return y_test, y_pred
+
+    # Compute the metrics for each validation fold in parallel.
+    jobs: list[Callable] = []
+    for (train_index, test_index) in fold.split(embeddings):
+      X_train, y_train = embeddings[train_index], labels[train_index]
+      X_train, y_train, sample_weights = self._setup_training(X_train, y_train, implicit_negatives)
+      X_test, y_test = embeddings[test_index], labels[test_index]
+      model = clone(self._model)
+      jobs.append(delayed(_fit_and_score)(model, X_train, y_train, sample_weights, X_test, y_test))
+    results = Parallel(n_jobs=-1)(jobs)
+
+    y_test = np.concatenate([y_test for y_test, _ in results], axis=0)
+    y_pred = np.concatenate([y_pred for _, y_pred in results], axis=0)
+    if len(set(y_test)) < 2:
+      return None, 0.5
+    roc_auc_val = roc_auc_score(y_test, y_pred)
+    precision, recall, thresholds = precision_recall_curve(y_test, y_pred)
+    numerator = (1 + F_BETA_WEIGHT**2) * precision * recall
+    denom = (F_BETA_WEIGHT**2 * precision) + recall
+    f1_scores = np.divide(numerator, denom, out=np.zeros_like(denom), where=(denom != 0))
+    max_f1: float = np.max(f1_scores)
+    max_f1_index = np.argmax(f1_scores)
+    max_f1_thresh: float = thresholds[max_f1_index]
+    max_f1_prec: float = precision[max_f1_index]
+    max_f1_recall: float = recall[max_f1_index]
+    metrics = ConceptMetrics(
+      f1=max_f1,
+      precision=max_f1_prec,
+      recall=max_f1_recall,
+      roc_auc=roc_auc_val,
+      overall=_get_overall_score(max_f1))
+    return metrics, max_f1_thresh
 
 
 def draft_examples(concept: Concept, draft: DraftId) -> dict[str, Example]:
@@ -136,7 +261,7 @@ def draft_examples(concept: Concept, draft: DraftId) -> dict[str, Example]:
     raise ValueError(
       f'Draft {draft} not found in concept. Found drafts: {list(draft_examples.keys())}')
 
-  # Map the text of the draft to its id so we can dedup with main.
+  # Map the text of the draft to its id so we can dedupe with main.
   draft_text_ids = {example.text: id for id, example in draft_examples[draft].items()}
 
   # Write each of examples from main to the draft examples only if the text does not appear in the
@@ -148,7 +273,8 @@ def draft_examples(concept: Concept, draft: DraftId) -> dict[str, Example]:
   return draft_examples[draft]
 
 
-class ConceptModel(BaseModel):
+@dataclasses.dataclass
+class ConceptModel:
   """A concept model. Stores all concept model drafts and manages syncing."""
   # The concept that this model is for.
   namespace: str
@@ -158,20 +284,27 @@ class ConceptModel(BaseModel):
   embedding_name: str
   version: int = -1
 
-  # The following fields are excluded from JSON serialization, but still pickleable.
+  column_info: Optional[ConceptColumnInfo] = None
+
+  # The following fields are excluded from JSON serialization, but still pickle-able.
   # Maps a concept id to the embeddings.
-  _embeddings: dict[str, np.ndarray] = {}
-  _logistic_models: dict[DraftId, LogisticEmbeddingModel] = {}
+  _embeddings: dict[str, np.ndarray] = dataclasses.field(default_factory=dict)
+  _logistic_models: dict[DraftId, LogisticEmbeddingModel] = dataclasses.field(default_factory=dict)
   _negative_vectors: Optional[np.ndarray] = None
 
-  class Config:
-    arbitrary_types_allowed = True
-    underscore_attrs_are_private = True
+  def get_metrics(self, concept: Concept) -> Optional[ConceptMetrics]:
+    """Return the metrics for this model."""
+    return self._get_logistic_model(DRAFT_MAIN)._metrics
+
+  def __post_init__(self) -> None:
+    if self.column_info:
+      self.column_info.path = normalize_path(self.column_info.path)
+      self._calibrate_on_dataset(self.column_info)
 
-  def calibrate_on_dataset(self, column_info: ConceptColumnInfo) -> None:
+  def _calibrate_on_dataset(self, column_info: ConceptColumnInfo) -> None:
     """Calibrate the model on the embeddings in the provided vector store."""
     db = get_dataset(column_info.namespace, column_info.name)
-    vector_store = db.get_vector_store(normalize_path(column_info.path))
+    vector_store = db.get_vector_store(self.embedding_name, normalize_path(column_info.path))
     keys = vector_store.keys()
     num_samples = min(column_info.num_negative_examples, len(keys))
     sample_keys = random.sample(keys, num_samples)
@@ -199,11 +332,7 @@ class ConceptModel(BaseModel):
   def _get_logistic_model(self, draft: DraftId) -> LogisticEmbeddingModel:
     """Get the logistic model for the provided draft."""
     if draft not in self._logistic_models:
-      self._logistic_models[draft] = LogisticEmbeddingModel(
-        namespace=self.namespace,
-        concept_name=self.concept_name,
-        embedding_name=self.embedding_name,
-        version=-1)
+      self._logistic_models[draft] = LogisticEmbeddingModel()
     return self._logistic_models[draft]
 
   def sync(self, concept: Concept) -> bool:
@@ -222,21 +351,9 @@ class ConceptModel(BaseModel):
       examples = draft_examples(concept, draft)
       embeddings = np.array([self._embeddings[id] for id in examples.keys()])
       labels = [example.label for example in examples.values()]
-      num_pos_labels = len([x for x in labels if x])
-      num_neg_labels = len([x for x in labels if not x])
-      sample_weights = [(1.0 / num_pos_labels if x else 1.0 / num_neg_labels) for x in labels]
-      if self._negative_vectors is not None:
-        num_implicit_labels = len(self._negative_vectors)
-        embeddings = np.concatenate([self._negative_vectors, embeddings])
-        labels = [False] * num_implicit_labels + labels
-        sample_weights = [1.0 / num_implicit_labels] * num_implicit_labels + sample_weights
-
       model = self._get_logistic_model(draft)
       with DebugTimer(f'Fitting model for "{concept_path}"'):
-        model.fit(embeddings, labels, sample_weights)
-
-      # Synchronize the model version with the concept version.
-      model.version = concept.version
+        model.fit(embeddings, labels, self._negative_vectors)
 
     # Synchronize the model version with the concept version.
     self.version = concept.version

src/concepts/db_concept.py

@@ -2,15 +2,18 @@
 
 import abc
 import glob
+import json
 import os
 import pickle
 import shutil
 
 # NOTE: We have to import the module for uuid so it can be mocked.
 import uuid
+from pathlib import Path
 from typing import List, Optional, Union, cast
 
 from pydantic import BaseModel
+from pyparsing import Any
 from typing_extensions import override
 
 from ..config import data_path
@@ -27,6 +30,7 @@ from .concept import (
   ExampleIn,
 )
 
+CONCEPTS_DIR = 'concept'
 DATASET_CONCEPTS_DIR = '.concepts'
 CONCEPT_JSON_FILENAME = 'concept.json'
 
@@ -65,8 +69,19 @@ class ConceptDB(abc.ABC):
     pass
 
   @abc.abstractmethod
-  def create(self, namespace: str, name: str, type: SignalInputType) -> Concept:
-    """Create a concept."""
+  def create(self,
+             namespace: str,
+             name: str,
+             type: SignalInputType,
+             description: Optional[str] = None) -> Concept:
+    """Create a concept.
+
+    Args:
+      namespace: The namespace of the concept.
+      name: The name of the concept.
+      type: The input type of the concept.
+      description: The description of the concept.
+    """
     pass
 
   @abc.abstractmethod
@@ -115,7 +130,7 @@ class ConceptModelDB(abc.ABC):
     pass
 
   @abc.abstractmethod
-  def _save(self, model: ConceptModel, column_info: Optional[ConceptColumnInfo]) -> None:
+  def _save(self, model: ConceptModel) -> None:
     """Save the concept model."""
     pass
 
@@ -126,13 +141,14 @@ class ConceptModelDB(abc.ABC):
       raise ValueError(f'Concept "{model.namespace}/{model.concept_name}" does not exist.')
     return concept.version == model.version
 
-  def sync(self, model: ConceptModel, column_info: Optional[ConceptColumnInfo]) -> bool:
+  def sync(self, model: ConceptModel) -> bool:
     """Sync the concept model. Returns true if the model was updated."""
     concept = self._concept_db.get(model.namespace, model.concept_name)
     if not concept:
       raise ValueError(f'Concept "{model.namespace}/{model.concept_name}" does not exist.')
     model_updated = model.sync(concept)
-    self._save(model, column_info)
+    if model_updated:
+      self._save(model)
     return model_updated
 
   @abc.abstractmethod
@@ -149,6 +165,16 @@ class ConceptModelDB(abc.ABC):
     """Remove all the models associated with a concept."""
     pass
 
+  @abc.abstractmethod
+  def get_models(self, namespace: str, concept_name: str) -> list[ConceptModel]:
+    """List all the models associated with a concept."""
+    pass
+
+  @abc.abstractmethod
+  def get_column_infos(self, namespace: str, concept_name: str) -> list[ConceptColumnInfo]:
+    """Get the dataset columns where this concept was applied to."""
+    pass
+
 
 class DiskConceptModelDB(ConceptModelDB):
   """Interface for the concept model database."""
@@ -191,10 +217,10 @@ class DiskConceptModelDB(ConceptModelDB):
     with open_file(concept_model_path, 'rb') as f:
       return pickle.load(f)
 
-  def _save(self, model: ConceptModel, column_info: Optional[ConceptColumnInfo]) -> None:
+  def _save(self, model: ConceptModel) -> None:
     """Save the concept model."""
     concept_model_path = _concept_model_path(model.namespace, model.concept_name,
-                                             model.embedding_name, column_info)
+                                             model.embedding_name, model.column_info)
     with open_file(concept_model_path, 'wb') as f:
       pickle.dump(model, f)
 
@@ -224,10 +250,39 @@ class DiskConceptModelDB(ConceptModelDB):
     for dir in dirs:
       shutil.rmtree(dir, ignore_errors=True)
 
+  @override
+  def get_models(self, namespace: str, concept_name: str) -> list[ConceptModel]:
+    """List all the models associated with a concept."""
+    model_files = glob.iglob(os.path.join(_concept_output_dir(namespace, concept_name), '*.pkl'))
+    models: list[ConceptModel] = []
+    for model_file in model_files:
+      embedding_name = os.path.basename(model_file)[:-len('.pkl')]
+      model = self.get(namespace, concept_name, embedding_name)
+      if model:
+        models.append(model)
+    return models
+
+  @override
+  def get_column_infos(self, namespace: str, concept_name: str) -> list[ConceptColumnInfo]:
+    datasets_path = os.path.join(data_path(), DATASETS_DIR_NAME)
+    # Skip if 'datasets' doesn't exist.
+    if not os.path.isdir(datasets_path):
+      return []
+
+    dirs = glob.iglob(
+      os.path.join(datasets_path, '**', DATASET_CONCEPTS_DIR, namespace, concept_name),
+      recursive=True)
+    result: list[ConceptColumnInfo] = []
+    for dir in dirs:
+      dir = os.path.relpath(dir, datasets_path)
+      dataset_namespace, dataset_name, *path, _, _, _ = Path(dir).parts
+      result.append(ConceptColumnInfo(namespace=dataset_namespace, name=dataset_name, path=path))
+    return result
+
 
 def _concept_output_dir(namespace: str, name: str) -> str:
   """Return the output directory for a given concept."""
-  return os.path.join(data_path(), 'concept', namespace, name)
+  return os.path.join(data_path(), CONCEPTS_DIR, namespace, name)
 
 
 def _concept_json_path(namespace: str, name: str) -> str:
@@ -246,7 +301,7 @@ def _concept_model_path(namespace: str,
   path_without_wildcards = (p for p in path_tuple if p != PATH_WILDCARD)
   path_dir = os.path.join(dataset_dir, *path_without_wildcards)
   return os.path.join(path_dir, DATASET_CONCEPTS_DIR, namespace, concept_name,
-                      f'{embedding_name}.pkl')
+                      f'{embedding_name}-neg-{column_info.num_negative_examples}.pkl')
 
 
 class DiskConceptDB(ConceptDB):
@@ -280,16 +335,29 @@ class DiskConceptDB(ConceptDB):
       return None
 
     with open_file(concept_json_path) as f:
-      return Concept.parse_raw(f.read())
+      obj: dict[str, Any] = json.load(f)
+      if 'namespace' not in obj:
+        obj['namespace'] = namespace
+      return Concept.parse_obj(obj)
 
   @override
-  def create(self, namespace: str, name: str, type: SignalInputType) -> Concept:
+  def create(self,
+             namespace: str,
+             name: str,
+             type: SignalInputType,
+             description: Optional[str] = None) -> Concept:
     """Create a concept."""
     concept_json_path = _concept_json_path(namespace, name)
     if file_exists(concept_json_path):
       raise ValueError(f'Concept with namespace "{namespace}" and name "{name}" already exists.')
 
-    concept = Concept(namespace=namespace, concept_name=name, type=type, data={}, version=0)
+    concept = Concept(
+      namespace=namespace,
+      concept_name=name,
+      type=type,
+      data={},
+      version=0,
+      description=description)
     self._save(concept)
     return concept
 

src/concepts/db_concept_test.py

@@ -1,7 +1,7 @@
 """Tests for the the database concept."""
 
 from pathlib import Path
-from typing import Generator, Iterable, Type, cast
+from typing import Generator, Iterable, Optional, Type, cast
 
 import numpy as np
 import pytest
@@ -423,7 +423,8 @@ class TestLogisticModel(LogisticEmbeddingModel):
     return np.array([.1])
 
   @override
-  def fit(self, embeddings: np.ndarray, labels: list[bool], sample_weights: list[float]) -> None:
+  def fit(self, embeddings: np.ndarray, labels: list[bool],
+          implicit_negatives: Optional[np.ndarray]) -> None:
     pass
 
 
@@ -436,20 +437,24 @@ class ConceptModelDBSuite:
     concept_db = concept_db_cls()
     model_db = model_db_cls(concept_db)
     model = _make_test_concept_model(concept_db)
-    model_db.sync(model, column_info=None)
+    model_db.sync(model)
     retrieved_model = model_db.get(
       namespace='test', concept_name='test_concept', embedding_name='test_embedding')
     if not retrieved_model:
       retrieved_model = model_db.create(
         namespace='test', concept_name='test_concept', embedding_name='test_embedding')
-    assert retrieved_model == model
+    assert retrieved_model.namespace == model.namespace
+    assert retrieved_model.concept_name == model.concept_name
+    assert retrieved_model.embedding_name == model.embedding_name
+    assert retrieved_model.version == model.version
+    assert retrieved_model.column_info == model.column_info
 
   def test_sync_model(self, concept_db_cls: Type[ConceptDB], model_db_cls: Type[ConceptModelDB],
                       mocker: MockerFixture) -> None:
 
     concept_db = concept_db_cls()
     model_db = model_db_cls(concept_db)
-    logistic_model = TestLogisticModel(embedding_name='test_embedding')
+    logistic_model = TestLogisticModel()
     score_embeddings_mock = mocker.spy(TestLogisticModel, 'score_embeddings')
     fit_mock = mocker.spy(TestLogisticModel, 'fit')
 
@@ -459,7 +464,7 @@ class ConceptModelDBSuite:
     assert score_embeddings_mock.call_count == 0
     assert fit_mock.call_count == 0
 
-    model_db.sync(model, column_info=None)
+    model_db.sync(model)
 
     assert model_db.in_sync(model) is True
     assert score_embeddings_mock.call_count == 0
@@ -471,20 +476,20 @@ class ConceptModelDBSuite:
     model_db = model_db_cls(concept_db)
     score_embeddings_mock = mocker.spy(TestLogisticModel, 'score_embeddings')
     fit_mock = mocker.spy(TestLogisticModel, 'fit')
-    logistic_model = TestLogisticModel(embedding_name='test_embedding')
+    logistic_model = TestLogisticModel()
     model = _make_test_concept_model(concept_db, logistic_models={DRAFT_MAIN: logistic_model})
-    model_db.sync(model, column_info=None)
+    model_db.sync(model)
     assert model_db.in_sync(model) is True
     assert score_embeddings_mock.call_count == 0
     assert fit_mock.call_count == 1
 
     (called_model, called_embeddings, called_labels,
-     called_weights) = fit_mock.call_args_list[-1].args
+     called_implicit_negatives) = fit_mock.call_args_list[-1].args
     assert called_model == logistic_model
     np.testing.assert_array_equal(
       called_embeddings, np.array([EMBEDDING_MAP['not in concept'], EMBEDDING_MAP['in concept']]))
     assert called_labels == [False, True]
-    assert called_weights == [1.0, 1.0]
+    assert called_implicit_negatives is None
 
     # Edit the concept.
     concept_db.edit('test', 'test_concept',
@@ -495,13 +500,13 @@ class ConceptModelDBSuite:
     assert score_embeddings_mock.call_count == 0
     assert fit_mock.call_count == 1
 
-    model_db.sync(model, column_info=None)
+    model_db.sync(model)
     assert model_db.in_sync(model) is True
     assert score_embeddings_mock.call_count == 0
     assert fit_mock.call_count == 2
     # Fit is called again with new points on main only.
     (called_model, called_embeddings, called_labels,
-     called_weights) = fit_mock.call_args_list[-1].args
+     called_implicit_negatives) = fit_mock.call_args_list[-1].args
     assert called_model == logistic_model
     np.testing.assert_array_equal(
       called_embeddings,
@@ -510,7 +515,7 @@ class ConceptModelDBSuite:
         EMBEDDING_MAP['a new data point']
       ]))
     assert called_labels == [False, True, False]
-    assert called_weights == pytest.approx([1 / 2, 1.0, 1 / 2])
+    assert called_implicit_negatives is None
 
   def test_out_of_sync_draft_model(self, concept_db_cls: Type[ConceptDB],
                                    model_db_cls: Type[ConceptModelDB],
@@ -519,14 +524,14 @@ class ConceptModelDBSuite:
     model_db = model_db_cls(concept_db)
     score_embeddings_mock = mocker.spy(TestLogisticModel, 'score_embeddings')
     fit_mock = mocker.spy(TestLogisticModel, 'fit')
-    logistic_model = TestLogisticModel(embedding_name='test_embedding')
-    draft_model = TestLogisticModel(embedding_name='test_embedding')
+    main_model = TestLogisticModel()
+    draft_model = TestLogisticModel()
     model = _make_test_concept_model(
       concept_db, logistic_models={
-        DRAFT_MAIN: logistic_model,
+        DRAFT_MAIN: main_model,
         'test_draft': draft_model
       })
-    model_db.sync(model, column_info=None)
+    model_db.sync(model)
     assert model_db.in_sync(model) is True
     assert score_embeddings_mock.call_count == 0
     assert fit_mock.call_count == 1
@@ -547,15 +552,16 @@ class ConceptModelDBSuite:
     assert score_embeddings_mock.call_count == 0
     assert fit_mock.call_count == 1
 
-    model_db.sync(model, column_info=None)
+    model_db.sync(model)
     assert model_db.in_sync(model) is True
     assert score_embeddings_mock.call_count == 0
     assert fit_mock.call_count == 3  # Fit is called on both the draft, and main.
 
     # Fit is called again with the same points.
-    ((called_model, called_embeddings, called_labels, called_weights),
-     (called_draft_model, called_draft_embeddings, called_draft_labels, called_draft_weights)) = (
-       c.args for c in fit_mock.call_args_list[-2:])
+    ((called_model, called_embeddings, called_labels, called_implicit_negatives),
+     (called_draft_model, called_draft_embeddings, called_draft_labels,
+      called_draft_implicit_negatives)) = (
+        c.args for c in fit_mock.call_args_list[-2:])
 
     # The draft model is called with the data from main, and the data from draft.
     assert called_draft_model == draft_model
@@ -572,21 +578,21 @@ class ConceptModelDBSuite:
       False,
       False
     ]
-    assert called_draft_weights == pytest.approx([1.0, 1 / 3, 1 / 3, 1 / 3])
+    assert called_draft_implicit_negatives is None
 
     # The main model was fit without the data from the draft.
-    assert called_model == draft_model
+    assert called_model == main_model
     np.testing.assert_array_equal(
       called_embeddings, np.array([EMBEDDING_MAP['not in concept'], EMBEDDING_MAP['in concept']]))
     assert called_labels == [False, True]
-    assert called_weights == pytest.approx([1.0, 1.0])
+    assert called_implicit_negatives is None
 
   def test_embedding_not_found_in_map(self, concept_db_cls: Type[ConceptDB],
                                       model_db_cls: Type[ConceptModelDB]) -> None:
     concept_db = concept_db_cls()
     model_db = model_db_cls(concept_db)
     model = _make_test_concept_model(concept_db)
-    model_db.sync(model, column_info=None)
+    model_db.sync(model)
 
     # Edit the concept.
     concept_db.edit('test', 'test_concept',
@@ -596,5 +602,5 @@ class ConceptModelDBSuite:
     assert model_db.in_sync(model) is False
 
     with pytest.raises(ValueError, match='Example "unknown text" not in embedding map'):
-      model_db.sync(model, column_info=None)
-      model_db.sync(model, column_info=None)
+      model_db.sync(model)
+      model_db.sync(model)

src/data/dataset.py

@@ -255,7 +255,7 @@ class Dataset(abc.ABC):
     pass
 
   @abc.abstractmethod
-  def get_vector_store(self, path: PathTuple) -> VectorStore:
+  def get_vector_store(self, embedding: str, path: PathTuple) -> VectorStore:
     # TODO: Instead of this, allow selecting vectors via select_rows.
     """Get the vector store for a column."""
     pass

src/data/dataset_duckdb.py

@@ -241,15 +241,23 @@ class DatasetDuckDB(Dataset):
     raise NotImplementedError('count is not yet implemented for DuckDB.')
 
   @override
-  def get_vector_store(self, path: PathTuple) -> VectorStore:
+  def get_vector_store(self, embedding: str, path: PathTuple) -> VectorStore:
     # Refresh the manifest to make sure we have the latest signal manifests.
     self.manifest()
 
+    if path[-1] != EMBEDDING_KEY:
+      path = (*path, embedding, PATH_WILDCARD, EMBEDDING_KEY)
+
     if path not in self._col_vector_stores:
-      manifest = next(
-        m for m in self._signal_manifests if schema_contains_path(m.data_schema, path))
-      if not manifest:
+      manifests = [
+        m for m in self._signal_manifests
+        if schema_contains_path(m.data_schema, path) and m.embedding_filename_prefix
+      ]
+      if not manifests:
         raise ValueError(f'No embedding found for path {path}.')
+      if len(manifests) > 1:
+        raise ValueError(f'Multiple embeddings found for path {path}. Got: {manifests}')
+      manifest = manifests[0]
       if not manifest.embedding_filename_prefix:
         raise ValueError(f'Signal manifest for path {path} is not an embedding. '
                          f'Got signal manifest: {manifest}')
@@ -273,7 +281,7 @@ class DatasetDuckDB(Dataset):
       manifest: DatasetManifest,
       compute_dependencies: Optional[bool] = False,
       task_step_id: Optional[TaskStepId] = None) -> tuple[PathTuple, Optional[TaskStepId]]:
-    """Run all the signals depedencies required to run this signal.
+    """Run all the signals dependencies required to run this signal.
 
     Args:
       signal: The signal to prepare.
@@ -560,7 +568,8 @@ class DatasetDuckDB(Dataset):
     if is_ordinal(leaf.dtype):
       min_max_query = f"""
         SELECT MIN(val) AS minVal, MAX(val) AS maxVal
-        FROM (SELECT {inner_select} as val FROM t);
+        FROM (SELECT {inner_select} as val FROM t)
+        WHERE NOT isnan(val)
       """
       row = self._query(min_max_query)[0]
       result.min_val, result.max_val = row
@@ -590,7 +599,9 @@ class DatasetDuckDB(Dataset):
     named_bins = _normalize_bins(bins or leaf.bins)
     stats = self.stats(leaf_path)
 
-    if is_float(leaf.dtype) or is_integer(leaf.dtype):
+    leaf_is_float = is_float(leaf.dtype)
+    leaf_is_integer = is_integer(leaf.dtype)
+    if leaf_is_float or leaf_is_integer:
       if named_bins is None:
         # Auto-bin.
         named_bins = _auto_bins(stats, NUM_AUTO_BINS)
@@ -606,11 +617,14 @@ class DatasetDuckDB(Dataset):
       bin_index_col = 'col0'
       bin_min_col = 'col1'
       bin_max_col = 'col2'
-      # We cast the field to `double` so bining works for both `float` and `int` fields.
+      is_nan_filter = f'NOT isnan({inner_val}) AND' if leaf_is_float else ''
+
+      # We cast the field to `double` so binning works for both `float` and `int` fields.
       outer_select = f"""(
         SELECT {bin_index_col} FROM (
           VALUES {', '.join(sql_bounds)}
-        ) WHERE {inner_val}::DOUBLE >= {bin_min_col} AND {inner_val}::DOUBLE < {bin_max_col}
+        ) WHERE {is_nan_filter}
+           {inner_val}::DOUBLE >= {bin_min_col} AND {inner_val}::DOUBLE < {bin_max_col}
       )"""
     else:
       if stats.approx_count_distinct >= dataset.TOO_MANY_DISTINCT:
@@ -625,6 +639,7 @@ class DatasetDuckDB(Dataset):
 
     filters, _ = self._normalize_filters(filters, col_aliases={}, udf_aliases={}, manifest=manifest)
     filter_queries = self._create_where(manifest, filters, searches=[])
+
     where_query = ''
     if filter_queries:
       where_query = f"WHERE {' AND '.join(filter_queries)}"
@@ -756,8 +771,9 @@ class DatasetDuckDB(Dataset):
     for udf_col in udf_columns:
       if isinstance(udf_col.signal_udf, ConceptScoreSignal):
         # Set dataset information on the signal.
+        source_path = udf_col.path if udf_col.path[-1] != EMBEDDING_KEY else udf_col.path[:-3]
         udf_col.signal_udf.set_column_info(
-          ConceptColumnInfo(namespace=self.namespace, name=self.dataset_name, path=udf_col.path))
+          ConceptColumnInfo(namespace=self.namespace, name=self.dataset_name, path=source_path))
 
     # Decide on the exact sorting order.
     sort_results = self._merge_sorts(search_udfs, sort_by, sort_order)
@@ -791,19 +807,19 @@ class DatasetDuckDB(Dataset):
     if topk_udf_col:
       key_prefixes: Optional[list[VectorKey]] = None
       if where_query:
-        # If there are filters, we need to send UUIDs to the topk query.
+        # If there are filters, we need to send UUIDs to the top k query.
         df = con.execute(f'SELECT {UUID_COLUMN} FROM t {where_query}').df()
         total_num_rows = len(df)
         key_prefixes = df[UUID_COLUMN]
 
-      signal = cast(Signal, topk_udf_col.signal_udf)
+      topk_signal = cast(TextEmbeddingModelSignal, topk_udf_col.signal_udf)
       # The input is an embedding.
-      vector_store = self.get_vector_store(topk_udf_col.path)
+      vector_store = self.get_vector_store(topk_signal.embedding, topk_udf_col.path)
       k = (limit or 0) + (offset or 0)
-      topk = signal.vector_compute_topk(k, vector_store, key_prefixes)
+      topk = topk_signal.vector_compute_topk(k, vector_store, key_prefixes)
       topk_uuids = list(dict.fromkeys([cast(str, key[0]) for key, _ in topk]))
 
-      # Ignore all the other filters and filter DuckDB results only by the topk UUIDs.
+      # Ignore all the other filters and filter DuckDB results only by the top k UUIDs.
       uuid_filter = Filter(path=(UUID_COLUMN,), op=ListOp.IN, value=topk_uuids)
       filter_query = self._create_where(manifest, [uuid_filter])[0]
       where_query = f'WHERE {filter_query}'
@@ -923,7 +939,8 @@ class DatasetDuckDB(Dataset):
 
         if signal.compute_type in [SignalInputType.TEXT_EMBEDDING]:
           # The input is an embedding.
-          vector_store = self.get_vector_store(udf_col.path)
+          embedding_signal = cast(TextEmbeddingModelSignal, signal)
+          vector_store = self.get_vector_store(embedding_signal.embedding, udf_col.path)
           flat_keys = flatten_keys(df[UUID_COLUMN], input)
           signal_out = signal.vector_compute(flat_keys, vector_store)
           # Add progress.
@@ -1273,39 +1290,48 @@ class DatasetDuckDB(Dataset):
     binary_ops = set(BinaryOp)
     unary_ops = set(UnaryOp)
     list_ops = set(ListOp)
-    for filter in filters:
-      duckdb_path = self._leaf_path_to_duckdb_path(filter.path, manifest.data_schema)
+    for f in filters:
+      duckdb_path = self._leaf_path_to_duckdb_path(f.path, manifest.data_schema)
       select_str = _select_sql(duckdb_path, flatten=True, unnest=False)
-      is_array = any(subpath == PATH_WILDCARD for subpath in filter.path)
+      is_array = any(subpath == PATH_WILDCARD for subpath in f.path)
+
+      nan_filter = ''
+      field = manifest.data_schema.get_field(f.path)
+      filter_nans = field.dtype and is_float(field.dtype)
 
-      if filter.op in binary_ops:
-        sql_op = BINARY_OP_TO_SQL[cast(BinaryOp, filter.op)]
-        filter_val = cast(FeatureValue, filter.value)
+      if f.op in binary_ops:
+        sql_op = BINARY_OP_TO_SQL[cast(BinaryOp, f.op)]
+        filter_val = cast(FeatureValue, f.value)
         if isinstance(filter_val, str):
           filter_val = f"'{filter_val}'"
         elif isinstance(filter_val, bytes):
           filter_val = _bytes_to_blob_literal(filter_val)
         else:
           filter_val = str(filter_val)
-        filter_query = (f'len(list_filter({select_str}, x -> x {sql_op} {filter_val})) > 0'
-                        if is_array else f'{select_str} {sql_op} {filter_val}')
-      elif filter.op in unary_ops:
-        if filter.op == UnaryOp.EXISTS:
+        if is_array:
+          nan_filter = 'NOT isnan(x) AND' if filter_nans else ''
+          filter_query = (f'len(list_filter({select_str}, '
+                          f'x -> {nan_filter} x {sql_op} {filter_val})) > 0')
+        else:
+          nan_filter = f'NOT isnan({select_str}) AND' if filter_nans else ''
+          filter_query = f'{nan_filter} {select_str} {sql_op} {filter_val}'
+      elif f.op in unary_ops:
+        if f.op == UnaryOp.EXISTS:
           filter_query = f'len({select_str}) > 0' if is_array else f'{select_str} IS NOT NULL'
         else:
-          raise ValueError(f'Unary op: {filter.op} is not yet supported')
-      elif filter.op in list_ops:
-        if filter.op == ListOp.IN:
-          filter_list_val = cast(FeatureListValue, filter.value)
+          raise ValueError(f'Unary op: {f.op} is not yet supported')
+      elif f.op in list_ops:
+        if f.op == ListOp.IN:
+          filter_list_val = cast(FeatureListValue, f.value)
           if not isinstance(filter_list_val, list):
             raise ValueError('filter with array value can only use the IN comparison')
           wrapped_filter_val = [f"'{part}'" for part in filter_list_val]
           filter_val = f'({", ".join(wrapped_filter_val)})'
           filter_query = f'{select_str} IN {filter_val}'
         else:
-          raise ValueError(f'List op: {filter.op} is not yet supported')
+          raise ValueError(f'List op: {f.op} is not yet supported')
       else:
-        raise ValueError(f'Invalid filter op: {filter.op}')
+        raise ValueError(f'Invalid filter op: {f.op}')
       sql_filter_queries.append(filter_query)
     return sql_filter_queries
 
@@ -1330,7 +1356,7 @@ class DatasetDuckDB(Dataset):
     return rows
 
   def _query_df(self, query: str) -> pd.DataFrame:
-    """Execute a query that returns a dataframe."""
+    """Execute a query that returns a data frame."""
     result = self._execute(query)
     df = _replace_nan_with_none(result.df())
     result.close()

src/data/dataset_select_groups_test.py

@@ -167,6 +167,8 @@ def test_named_bins(make_test_data: TestDataMaker) -> None:
     'age': 80
   }, {
     'age': 55
+  }, {
+    'age': float('nan')
   }]
   dataset = make_test_data(items)
 
@@ -178,7 +180,7 @@ def test_named_bins(make_test_data: TestDataMaker) -> None:
       ('middle-aged', 50, 65),
       ('senior', 65, None),
     ])
-  assert result.counts == [('adult', 2), ('young', 1), ('senior', 1), ('middle-aged', 1)]
+  assert result.counts == [('adult', 2), ('young', 1), ('senior', 1), ('middle-aged', 1), (None, 1)]
 
 
 def test_schema_with_bins(make_test_data: TestDataMaker) -> None:
@@ -192,11 +194,13 @@ def test_schema_with_bins(make_test_data: TestDataMaker) -> None:
     'age': 80
   }, {
     'age': 55
+  }, {
+    'age': float('nan')
   }]
   data_schema = schema({
     UUID_COLUMN: 'string',
     'age': field(
-      'int32',
+      'float32',
       bins=[
         ('young', None, 20),
         ('adult', 20, 50),
@@ -207,7 +211,7 @@ def test_schema_with_bins(make_test_data: TestDataMaker) -> None:
   dataset = make_test_data(items, data_schema)
 
   result = dataset.select_groups(leaf_path='age')
-  assert result.counts == [('adult', 2), ('young', 1), ('senior', 1), ('middle-aged', 1)]
+  assert result.counts == [('adult', 2), ('young', 1), ('senior', 1), ('middle-aged', 1), (None, 1)]
 
 
 def test_filters(make_test_data: TestDataMaker) -> None:
@@ -304,10 +308,10 @@ def test_too_many_distinct(make_test_data: TestDataMaker, mocker: MockerFixture)
 
 
 def test_auto_bins_for_float(make_test_data: TestDataMaker) -> None:
-  items: list[Item] = [{'feature': float(i)} for i in range(5)]
+  items: list[Item] = [{'feature': float(i)} for i in range(5)] + [{'feature': float('nan')}]
   dataset = make_test_data(items)
 
   res = dataset.select_groups('feature')
-  assert res.counts == [('0', 1), ('3', 1), ('7', 1), ('11', 1), ('14', 1)]
+  assert res.counts == [('0', 1), ('3', 1), ('7', 1), ('11', 1), ('14', 1), (None, 1)]
   assert res.too_many_distinct is False
   assert res.bins

src/data/dataset_select_rows_filter_test.py

@@ -24,6 +24,9 @@ TEST_DATA: list[Item] = [{
   'int': 2,
   'bool': True,
   'float': 1.0
+}, {
+  UUID_COLUMN: '4',
+  'float': float('nan')
 }]
 
 
@@ -46,6 +49,91 @@ def test_filter_by_ids(make_test_data: TestDataMaker) -> None:
   assert list(result) == []
 
 
+def test_filter_greater(make_test_data: TestDataMaker) -> None:
+  dataset = make_test_data(TEST_DATA)
+
+  id_filter: BinaryFilterTuple = ('float', BinaryOp.GREATER, 2.0)
+  result = dataset.select_rows(filters=[id_filter])
+
+  assert list(result) == [{UUID_COLUMN: '1', 'str': 'a', 'int': 1, 'bool': False, 'float': 3.0}]
+
+
+def test_filter_greater_equal(make_test_data: TestDataMaker) -> None:
+  dataset = make_test_data(TEST_DATA)
+
+  id_filter: BinaryFilterTuple = ('float', BinaryOp.GREATER_EQUAL, 2.0)
+  result = dataset.select_rows(filters=[id_filter])
+
+  assert list(result) == [{
+    UUID_COLUMN: '1',
+    'str': 'a',
+    'int': 1,
+    'bool': False,
+    'float': 3.0
+  }, {
+    UUID_COLUMN: '2',
+    'str': 'b',
+    'int': 2,
+    'bool': True,
+    'float': 2.0
+  }]
+
+
+def test_filter_less(make_test_data: TestDataMaker) -> None:
+  dataset = make_test_data(TEST_DATA)
+
+  id_filter: BinaryFilterTuple = ('float', BinaryOp.LESS, 2.0)
+  result = dataset.select_rows(filters=[id_filter])
+
+  assert list(result) == [{UUID_COLUMN: '3', 'str': 'b', 'int': 2, 'bool': True, 'float': 1.0}]
+
+
+def test_filter_less_equal(make_test_data: TestDataMaker) -> None:
+  dataset = make_test_data(TEST_DATA)
+
+  id_filter: BinaryFilterTuple = ('float', BinaryOp.LESS_EQUAL, 2.0)
+  result = dataset.select_rows(filters=[id_filter])
+
+  assert list(result) == [{
+    UUID_COLUMN: '2',
+    'str': 'b',
+    'int': 2,
+    'bool': True,
+    'float': 2.0
+  }, {
+    UUID_COLUMN: '3',
+    'str': 'b',
+    'int': 2,
+    'bool': True,
+    'float': 1.0
+  }]
+
+
+def test_filter_not_equal(make_test_data: TestDataMaker) -> None:
+  dataset = make_test_data(TEST_DATA)
+
+  id_filter: BinaryFilterTuple = ('float', BinaryOp.NOT_EQUAL, 2.0)
+  result = dataset.select_rows(filters=[id_filter])
+
+  assert list(result) == [
+    {
+      UUID_COLUMN: '1',
+      'str': 'a',
+      'int': 1,
+      'bool': False,
+      'float': 3.0
+    },
+    {
+      UUID_COLUMN: '3',
+      'str': 'b',
+      'int': 2,
+      'bool': True,
+      'float': 1.0
+    },
+    # NaNs are not counted when we are filtering a field.
+  ]
+
+
 def test_filter_by_list_of_ids(make_test_data: TestDataMaker) -> None:
   dataset = make_test_data(TEST_DATA)
 

src/data/dataset_select_rows_search_test.py

@@ -288,17 +288,9 @@ def test_concept_search(make_test_data: TestDataMaker, mocker: MockerFixture) ->
     },
   ]
 
-  # Make sure fit was called with negative examples.
   (_, embeddings, labels, _) = concept_model_mock.call_args_list[-1].args
-  assert embeddings.shape == (8, 3)
+  assert embeddings.shape == (2, 3)
   assert labels == [
-    # Negative implicit labels.
-    False,
-    False,
-    False,
-    False,
-    False,
-    False,
     # Explicit labels.
     False,
     True

src/data/dataset_stats_test.py

@@ -15,7 +15,7 @@ SIMPLE_ITEMS: list[Item] = [{
   'str': 'a',
   'int': 1,
   'bool': False,
-  'float': 3.0
+  'float': 3.0,
 }, {
   UUID_COLUMN: '2',
   'str': 'b',
@@ -28,6 +28,9 @@ SIMPLE_ITEMS: list[Item] = [{
   'int': 2,
   'bool': True,
   'float': 1.0
+}, {
+  UUID_COLUMN: '4',
+  'float': float('nan')
 }]
 
 
@@ -40,7 +43,7 @@ def test_simple_stats(make_test_data: TestDataMaker) -> None:
 
   result = dataset.stats(leaf_path='float')
   assert result == StatsResult(
-    path=('float',), total_count=3, approx_count_distinct=3, min_val=1.0, max_val=3.0)
+    path=('float',), total_count=4, approx_count_distinct=4, min_val=1.0, max_val=3.0)
 
   result = dataset.stats(leaf_path='bool')
   assert result == StatsResult(path=('bool',), total_count=3, approx_count_distinct=2)

src/data/dataset_utils.py

@@ -1,5 +1,6 @@
 """Utilities for working with datasets."""
 
+import json
 import math
 import os
 import pickle
@@ -283,7 +284,10 @@ def write_items_to_parquet(items: Iterable[Item], output_dir: str, schema: Schem
     if UUID_COLUMN not in item:
       item[UUID_COLUMN] = secrets.token_urlsafe(nbytes=12)  # 16 base64 characters.
     if os.getenv('DEBUG'):
-      _validate(item, arrow_schema)
+      try:
+        _validate(item, arrow_schema)
+      except Exception as e:
+        raise ValueError(f'Error validating item: {json.dumps(item)}') from e
     writer.write(item)
     num_items += 1
   writer.close()