diff --git a/internal/cli/dataset.go b/internal/cli/dataset.go
index 1b63d92d..8cb40baf 100644
--- a/internal/cli/dataset.go
+++ b/internal/cli/dataset.go
@@ -222,8 +222,7 @@ Exit codes:
cmd.Flags().StringVar(&table, "table", "",
"destination table name (MySQL identifier; matches /data/shared/
/ on the PVC)")
cmd.Flags().StringVar(&category, "category", "image_classification",
- "task category: image_classification, tabular_classification, tabular_regression, "+
- "time_series_forecasting, time_to_event_prediction")
+ "task category, one of: "+push.SupportedCategoriesList())
cmd.Flags().StringVar(&intent, "intent", "",
"intent: train|test")
cmd.Flags().StringVar(&labelColumn, "label-column", "",
@@ -422,24 +421,20 @@ contributors train against it without ever seeing the raw files.`))
case a.Spec.Category == "":
// Left empty by a caller; let the schema produce the canonical
// "category is required" error downstream.
- case push.IsTabular(a.Spec.Category) || push.IsText(a.Spec.Category) ||
- a.Spec.Category == "image_classification" ||
- a.Spec.Category == "object_detection" ||
- a.Spec.Category == "keypoint_detection":
+ case push.IsCLISupported(a.Spec.Category):
// supported
case push.IsImage(a.Spec.Category):
- // semantic_segmentation / instance_segmentation
+ // A known image category dataset push doesn't implement yet
+ // (semantic_segmentation / instance_segmentation). The per-category
+ // reason + the supported list both come from the registry.
+ spec, _ := push.Lookup(a.Spec.Category)
return &exitError{code: 2, err: fmt.Errorf(
- "category %q isn't supported by the CLI yet. semantic_segmentation is "+
- "blocked on the ingestor's mask-sidecar support (data-ingestors#136), and "+
- "instance_segmentation isn't implemented. Supported image categories: "+
- "image_classification, object_detection, keypoint_detection.", a.Spec.Category)}
+ "category %q isn't supported by the CLI yet (%s). Supported categories: %s.",
+ a.Spec.Category, spec.UnsupportedNote, push.SupportedCategoriesList())}
default:
return &exitError{code: 2, err: fmt.Errorf(
- "category %q isn't a recognized task category. Supported: image_classification, "+
- "object_detection, keypoint_detection, text_classification, "+
- "masked_language_modeling, tabular_classification, tabular_regression, "+
- "time_series_forecasting, time_to_event_prediction.", a.Spec.Category)}
+ "category %q isn't a recognized task category. Supported categories: %s.",
+ a.Spec.Category, push.SupportedCategoriesList())}
}
// 3. Walk the local directory FIRST (local "fail fast"), dispatched
diff --git a/internal/cli/interactive.go b/internal/cli/interactive.go
index 3c567bfd..bb6478d9 100644
--- a/internal/cli/interactive.go
+++ b/internal/cli/interactive.go
@@ -16,20 +16,12 @@ import (
)
// promptCategories is the ordered list offered by the interactive
-// category picker — the categories `dataset push` supports today (the
-// same set runDatasetPush's category gate accepts). semantic_ /
-// instance_segmentation are omitted until they're implemented.
-var promptCategories = []string{
- "image_classification",
- "object_detection",
- "keypoint_detection",
- "text_classification",
- "masked_language_modeling",
- "tabular_classification",
- "tabular_regression",
- "time_series_forecasting",
- "time_to_event_prediction",
-}
+// category picker. It derives from the push registry's CLI-supported
+// set — the exact categories runDatasetPush's gate accepts — so the
+// picker can't drift from what `dataset push` actually supports.
+// semantic_/instance_segmentation are excluded (CLISupported=false)
+// until they're implemented.
+var promptCategories = push.SupportedCategoryIDs()
// prompter is the narrow seam over the interactive library. Production
// uses surveyPrompter (a real terminal); tests inject a fake that
diff --git a/internal/push/category.go b/internal/push/category.go
index c20728a4..9ea14c3d 100644
--- a/internal/push/category.go
+++ b/internal/push/category.go
@@ -1,72 +1,152 @@
package push
-// Category families. These mirror data-ingestors'
+import "strings"
+
+// CategorySpec is the single source of truth for one task category's
+// CLI-relevant rules. It mirrors data-ingestors'
// tracebloc_ingestor/cli/conventions.py groupings so the CLI's
-// per-category behaviour (which flags are required, which local
-// layout to expect, which spec fields to emit) stays in lock-step
+// per-category behaviour (which local layout to expect, which spec
+// fields to emit, whether a label policy is needed) stays in lock-step
// with what the ingestor actually resolves.
//
-// Kept as a single source of truth here rather than scattered
-// string comparisons across spec.go / dataset.go.
-
-// imageCategories take a labels CSV + an images/ directory (plus,
-// for some, extra sidecar dirs handled in later increments).
-var imageCategories = map[string]bool{
- "image_classification": true,
- "object_detection": true,
- "keypoint_detection": true,
- "semantic_segmentation": true,
- "instance_segmentation": true,
+// Everything category-shaped derives from the registry below — the
+// family predicates, the `--category` help text, the interactive
+// picker, and the push accept-gate — so the enumerations can't drift
+// apart (they used to: the flag help listed 5 of 9, cli#74).
+type CategorySpec struct {
+ // ID is the canonical category identifier; it matches the
+ // ingest.v1 schema enum (vendored via scripts/sync-schema.sh).
+ ID string
+ // Family selects the local layout + staging shape.
+ Family Family
+ // Label is the human-friendly name shown in the interactive picker.
+ Label string
+ // RegressionClass marks categories that predict a numeric target and
+ // therefore need label.policy (object label form) so the raw target
+ // never ships to the central backend by default.
+ RegressionClass bool
+ // CLISupported reports whether `dataset push` implements the category
+ // today. semantic_/instance_segmentation are known (the schema
+ // defines them) but not yet pushable.
+ CLISupported bool
+ // UnsupportedNote explains why a known-but-unimplemented category
+ // isn't available yet; surfaced by the push gate. Empty when supported.
+ UnsupportedNote string
+}
+
+// Family groups categories by local layout.
+type Family int
+
+const (
+ // FamilyImage: a labels CSV + an images/ directory (plus, for some,
+ // extra sidecar dirs like annotations/ or masks/).
+ FamilyImage Family = iota
+ // FamilyTabular: a single CSV whose columns are described by a
+ // `schema` (column → SQL type) map. No sidecar files.
+ FamilyTabular
+ // FamilyText: a labels CSV + a directory of text files (texts/ for
+ // classification, sequences/ for masked language modeling).
+ FamilyText
+)
+
+// categoryRegistry is the ordered, authoritative list of every category
+// the ingest.v1 schema defines. Order is the display order for help text
+// and the interactive picker (CLI-supported first, in workflow order;
+// the not-yet-implemented ones last). Adding a category to the schema
+// means adding it here — the parity test pins the set.
+var categoryRegistry = []CategorySpec{
+ {ID: "image_classification", Family: FamilyImage, Label: "Image classification", CLISupported: true},
+ {ID: "object_detection", Family: FamilyImage, Label: "Object detection", CLISupported: true},
+ {ID: "keypoint_detection", Family: FamilyImage, Label: "Keypoint detection", CLISupported: true},
+ {ID: "text_classification", Family: FamilyText, Label: "Text classification", CLISupported: true},
+ {ID: "masked_language_modeling", Family: FamilyText, Label: "Masked language modeling", CLISupported: true},
+ {ID: "tabular_classification", Family: FamilyTabular, Label: "Tabular classification", CLISupported: true},
+ {ID: "tabular_regression", Family: FamilyTabular, Label: "Tabular regression", RegressionClass: true, CLISupported: true},
+ {ID: "time_series_forecasting", Family: FamilyTabular, Label: "Time-series forecasting", RegressionClass: true, CLISupported: true},
+ {ID: "time_to_event_prediction", Family: FamilyTabular, Label: "Time-to-event prediction", RegressionClass: true, CLISupported: true},
+ {ID: "semantic_segmentation", Family: FamilyImage, Label: "Semantic segmentation", CLISupported: false,
+ UnsupportedNote: "blocked on the ingestor's mask-sidecar support (data-ingestors#136)"},
+ {ID: "instance_segmentation", Family: FamilyImage, Label: "Instance segmentation", CLISupported: false,
+ UnsupportedNote: "not implemented"},
+}
+
+// categoryByID indexes the registry for O(1) lookup, built once.
+var categoryByID = func() map[string]CategorySpec {
+ m := make(map[string]CategorySpec, len(categoryRegistry))
+ for _, c := range categoryRegistry {
+ m[c.ID] = c
+ }
+ return m
+}()
+
+// Lookup returns the spec for a category id and whether it is known.
+func Lookup(category string) (CategorySpec, bool) {
+ c, ok := categoryByID[category]
+ return c, ok
+}
+
+// IsKnown reports whether category is a recognized task category (in the
+// schema), supported by the CLI or not.
+func IsKnown(category string) bool {
+ _, ok := categoryByID[category]
+ return ok
}
-// tabularCategories take a single CSV whose columns are described by
-// a `schema` (column → SQL type) map. No sidecar files.
-var tabularCategories = map[string]bool{
- "tabular_classification": true,
- "tabular_regression": true,
- "time_series_forecasting": true,
- "time_to_event_prediction": true,
+// IsCLISupported reports whether `dataset push` implements category today.
+func IsCLISupported(category string) bool { return categoryByID[category].CLISupported }
+
+// IsImage reports whether category uses the labels.csv + images/ layout.
+func IsImage(category string) bool {
+ c, ok := categoryByID[category]
+ return ok && c.Family == FamilyImage
}
-// regressionClassCategories predict a numeric target rather than a
-// class. The schema requires the label in object form with an
-// explicit `policy` so the raw target never ships to the central
-// backend by default (policy=bucket bins it first).
-var regressionClassCategories = map[string]bool{
- "tabular_regression": true,
- "time_series_forecasting": true,
- "time_to_event_prediction": true,
+// IsTabular reports whether category uses the single-CSV + schema layout.
+func IsTabular(category string) bool {
+ c, ok := categoryByID[category]
+ return ok && c.Family == FamilyTabular
}
-// textCategories take a labels CSV + a directory of text files
-// (texts/ for classification, sequences/ for masked language
-// modeling). masked_language_modeling additionally needs a
-// tokenizer.json at the dataset root and has NO label.
-var textCategories = map[string]bool{
- "text_classification": true,
- "masked_language_modeling": true,
+// IsText reports whether category uses the labels.csv + text-file dir layout.
+func IsText(category string) bool {
+ c, ok := categoryByID[category]
+ return ok && c.Family == FamilyText
}
-// IsImage reports whether category uses the labels.csv + images/
-// local layout.
-func IsImage(category string) bool { return imageCategories[category] }
+// IsRegressionClass reports whether category predicts a numeric target and
+// therefore needs label.policy (object label form).
+func IsRegressionClass(category string) bool { return categoryByID[category].RegressionClass }
-// IsTabular reports whether category uses the single-CSV + schema
-// local layout (no sidecar files).
-func IsTabular(category string) bool { return tabularCategories[category] }
+// SupportedCategoryIDs returns the ids `dataset push` supports, in display
+// order. Used to build the --category help, the interactive picker, and
+// the accept-gate's "Supported:" lists from one place.
+func SupportedCategoryIDs() []string {
+ ids := make([]string, 0, len(categoryRegistry))
+ for _, c := range categoryRegistry {
+ if c.CLISupported {
+ ids = append(ids, c.ID)
+ }
+ }
+ return ids
+}
-// IsRegressionClass reports whether category predicts a numeric
-// target and therefore needs label.policy (object label form).
-func IsRegressionClass(category string) bool { return regressionClassCategories[category] }
+// AllCategoryIDs returns every recognized category id, in registry order.
+func AllCategoryIDs() []string {
+ ids := make([]string, 0, len(categoryRegistry))
+ for _, c := range categoryRegistry {
+ ids = append(ids, c.ID)
+ }
+ return ids
+}
-// IsText reports whether category uses the labels.csv + text-file
-// directory (texts/ or sequences/) local layout.
-func IsText(category string) bool { return textCategories[category] }
+// SupportedCategoriesList is the comma-joined supported ids, for help text
+// and gate error messages.
+func SupportedCategoriesList() string { return strings.Join(SupportedCategoryIDs(), ", ") }
// TextSidecarDir returns the sidecar directory name a text category
// expects: "sequences" for masked_language_modeling, "texts" for
-// text_classification. (Used both as the local subdir to stage and
-// the spec field to emit.)
+// text_classification. (Used both as the local subdir to stage and the
+// spec field to emit.)
func TextSidecarDir(category string) string {
if category == "masked_language_modeling" {
return "sequences"
diff --git a/internal/push/category_registry_test.go b/internal/push/category_registry_test.go
new file mode 100644
index 00000000..1bb40ea3
--- /dev/null
+++ b/internal/push/category_registry_test.go
@@ -0,0 +1,99 @@
+package push
+
+import (
+ "sort"
+ "testing"
+)
+
+// The registry is the single source of truth; these pin its contents and
+// that the family predicates + the supported set all derive from it, so a
+// future edit can't reintroduce the "5 of 9" drift (cli#74).
+
+func TestRegistryKnownCategories(t *testing.T) {
+ want := []string{
+ "image_classification", "object_detection", "keypoint_detection",
+ "semantic_segmentation", "instance_segmentation",
+ "text_classification", "masked_language_modeling",
+ "tabular_classification", "tabular_regression",
+ "time_series_forecasting", "time_to_event_prediction",
+ }
+ if got := AllCategoryIDs(); !equalSet(got, want) {
+ t.Fatalf("AllCategoryIDs() = %v, want set %v", got, want)
+ }
+ for _, id := range want {
+ if !IsKnown(id) {
+ t.Errorf("IsKnown(%q) = false, want true", id)
+ }
+ }
+ if IsKnown("not_a_category") {
+ t.Error(`IsKnown("not_a_category") = true, want false`)
+ }
+}
+
+func TestSupportedCategories(t *testing.T) {
+ got := SupportedCategoryIDs()
+ if len(got) != 9 {
+ t.Fatalf("SupportedCategoryIDs() len = %d, want 9: %v", len(got), got)
+ }
+ for _, id := range got {
+ if !IsCLISupported(id) {
+ t.Errorf("SupportedCategoryIDs returned %q but IsCLISupported is false", id)
+ }
+ }
+ // semantic_/instance_segmentation are known but not yet pushable, and
+ // must explain why.
+ for _, id := range []string{"semantic_segmentation", "instance_segmentation"} {
+ if !IsKnown(id) {
+ t.Errorf("%s should be known", id)
+ }
+ if IsCLISupported(id) {
+ t.Errorf("%s should not be CLI-supported yet", id)
+ }
+ if spec, _ := Lookup(id); spec.UnsupportedNote == "" {
+ t.Errorf("%s should carry an UnsupportedNote", id)
+ }
+ }
+}
+
+func TestPredicatesDeriveFromRegistry(t *testing.T) {
+ for _, c := range categoryRegistry {
+ switch c.Family {
+ case FamilyImage:
+ if !IsImage(c.ID) || IsTabular(c.ID) || IsText(c.ID) {
+ t.Errorf("%s: predicates disagree with FamilyImage", c.ID)
+ }
+ case FamilyTabular:
+ if !IsTabular(c.ID) || IsImage(c.ID) || IsText(c.ID) {
+ t.Errorf("%s: predicates disagree with FamilyTabular", c.ID)
+ }
+ case FamilyText:
+ if !IsText(c.ID) || IsImage(c.ID) || IsTabular(c.ID) {
+ t.Errorf("%s: predicates disagree with FamilyText", c.ID)
+ }
+ }
+ if IsRegressionClass(c.ID) != c.RegressionClass {
+ t.Errorf("%s: IsRegressionClass = %v, want %v", c.ID, IsRegressionClass(c.ID), c.RegressionClass)
+ }
+ }
+ // An unknown category: every predicate false (no panic on missing key).
+ if IsImage("nope") || IsTabular("nope") || IsText("nope") ||
+ IsRegressionClass("nope") || IsCLISupported("nope") {
+ t.Error("predicates should all be false for an unknown category")
+ }
+}
+
+func equalSet(a, b []string) bool {
+ if len(a) != len(b) {
+ return false
+ }
+ as := append([]string(nil), a...)
+ bs := append([]string(nil), b...)
+ sort.Strings(as)
+ sort.Strings(bs)
+ for i := range as {
+ if as[i] != bs[i] {
+ return false
+ }
+ }
+ return true
+}