Upgrade training interactions

tests/general/test_signals.py

@@ -641,6 +641,43 @@ def test_save_signals_with_dict_predictions(self, mock_gm, mock_get_dataframe):
 
             self.assertTrue(mock_df.enqueue_batch.called)
 
+    @patch("weightslab.src.get_dataframe")
+    @patch("weightslab.src._gm")
+    def test_save_signals_with_string_predictions(self, mock_gm, mock_get_dataframe):
+        """Test save_signals with TEXT predictions/labels (e.g. a generative
+        model's reply and its reference/target text). Must not crash (the
+        pre-fix behavior raised ValueError from the classification-label
+        uint16 cast on non-numeric strings) and must forward native Python
+        str values, unwrapped, so downstream RecordMetadata/DataStat
+        construction sees a real str rather than a 0-d numpy array."""
+        mock_df = MagicMock()
+        mock_get_dataframe.return_value = mock_df
+
+        mock_model = MagicMock()
+        mock_model.current_step = 1
+        mock_gm.return_value = mock_model
+
+        with patch("weightslab.src.DATAFRAME_M", mock_df):
+            batch_ids = torch.tensor([1, 2])
+            signals = {"loss": torch.tensor(0.5)}
+            preds = ["Certainly! Here is the answer to your question.", "I don't know."]
+            targets = ["The expected reference answer.", "Another reference answer."]
+
+            wl.save_signals(
+                signals=signals,
+                batch_ids=batch_ids,
+                preds=preds,
+                targets=targets,
+                log=True
+            )
+
+            self.assertTrue(mock_df.enqueue_batch.called)
+            _, kwargs = mock_df.enqueue_batch.call_args
+            self.assertEqual(kwargs["preds"], preds)
+            self.assertEqual(kwargs["targets"], targets)
+            for value in kwargs["preds"] + kwargs["targets"]:
+                self.assertIsInstance(value, str)
+
     # =========================================================================
     # Signal Composition and Combination Tests
     # =========================================================================

tests/trainer/services/test_data_service_sample_id_query.py

@@ -0,0 +1,75 @@
+"""Regression test: filtering by a numeric-looking sample_id (e.g. the UI's
+"sample_id == 11448" quick filter) must actually match, even though
+sample_id is always stored as a STRING column (see save_signals's
+`batch_ids_np = [str(i) for i in batch_ids]`).
+
+`df.query("sample_id == 11448")` does NOT raise -- pandas happily compares
+the string column to the int literal and evaluates to False for every row,
+so the existing df.eval/`_mask_from_coerced_query` fallback (which only
+triggers on an exception) never kicks in, and the UI reports "0 of N
+samples" for an id that genuinely exists.
+"""
+
+import unittest
+
+import pandas as pd
+
+from weightslab.trainer.services.data_service import DataService
+
+
+class _StubDataService:
+    """Minimal stand-in -- _apply_agent_operation's df.query branch only
+    calls the (static) _mask_from_coerced_query helper, no other self state.
+    Bound explicitly since a bare object has no relationship to DataService
+    for `self._mask_from_coerced_query` to resolve through."""
+    _mask_from_coerced_query = staticmethod(DataService._mask_from_coerced_query)
+
+
+class TestSampleIdNumericQuery(unittest.TestCase):
+    def setUp(self):
+        self.service = _StubDataService()
+
+    def _apply(self, df, expr):
+        return DataService._apply_agent_operation(
+            self.service, df, "df.query", {"expr": expr})
+
+    def test_numeric_literal_matches_a_string_sample_id_column(self):
+        df = pd.DataFrame({
+            "sample_id": ["11448", "22", "34938"],
+            "reward": [1.0, 2.0, 3.0],
+        })
+        msg = self._apply(df, "sample_id == 11448")
+        self.assertIn("Applied query", msg)
+        self.assertEqual(list(df["sample_id"]), ["11448"])
+
+    def test_still_works_for_genuinely_numeric_columns(self):
+        df = pd.DataFrame({
+            "sample_id": ["1", "2", "3"],
+            "reward": [1.0, 2.0, 3.0],
+        })
+        self._apply(df, "reward > 1.5")
+        self.assertEqual(list(df["sample_id"]), ["2", "3"])
+
+    def test_genuinely_empty_result_stays_empty(self):
+        df = pd.DataFrame({
+            "sample_id": ["1", "2", "3"],
+            "reward": [1.0, 2.0, 3.0],
+        })
+        self._apply(df, "reward > 1000")
+        self.assertEqual(len(df), 0)
+
+    def test_text_prediction_column_is_not_coerced_away(self):
+        # A string-typed 'prediction' column with genuine free text must not
+        # be affected by the numeric-coercion fallback -- confirms the fix is
+        # scoped to numeric-looking columns only, not text predictions/labels.
+        df = pd.DataFrame({
+            "sample_id": ["1", "2"],
+            "prediction": ["Certainly! Let's flip a coin.", "I don't know."],
+        })
+        msg = self._apply(df, "prediction == \"I don't know.\"")
+        self.assertIn("Applied query", msg)
+        self.assertEqual(list(df["sample_id"]), ["2"])
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/trainer/services/test_data_service_text_predictions.py

@@ -0,0 +1,148 @@
+"""Unit tests for TEXT label/prediction support in data_service.py.
+
+Covers two things added for generative tasks (e.g. an LLM's generated reply
+as "prediction", a reference/target string as "label"):
+
+1. `looks_like_file_path_label` -- the heuristic that decides whether a
+   string label is "empty" (a segmentation mask path with nothing loaded
+   yet) vs. real text that must be preserved as-is.
+2. `DataService._process_sample_row`'s classification/tabular branches,
+   which must emit a string DataStat (name='target'/'pred', type='string',
+   value_string=...) for text values instead of crashing on float(label)/
+   float(pred) and silently dropping the stat.
+"""
+
+import unittest
+
+from weightslab.data.sample_stats import SampleStatsEx
+from weightslab.proto import experiment_service_pb2 as pb2
+from weightslab.trainer.services.data_service import (
+    DataService,
+    looks_like_file_path_label,
+)
+
+
+class TestLooksLikeFilePathLabel(unittest.TestCase):
+    def test_true_for_segmentation_mask_paths(self):
+        self.assertTrue(looks_like_file_path_label("mask.png"))
+        self.assertTrue(looks_like_file_path_label("/data/masks/sample_001.jpg"))
+        self.assertTrue(looks_like_file_path_label("label.tif"))
+
+    def test_false_for_free_text_without_trailing_period(self):
+        # The exact regression this heuristic used to misfire on: ordinary
+        # multi-sentence text that doesn't end with a period.
+        self.assertFalse(looks_like_file_path_label(
+            "Certainly! Let's break this down. Here is more info"))
+        self.assertFalse(looks_like_file_path_label("See section 2.5 for details"))
+
+    def test_false_for_free_text_with_trailing_period(self):
+        self.assertFalse(looks_like_file_path_label("The reference answer."))
+
+    def test_false_for_non_string_or_no_period(self):
+        self.assertFalse(looks_like_file_path_label(None))
+        self.assertFalse(looks_like_file_path_label(42))
+        self.assertFalse(looks_like_file_path_label("no period here"))
+
+
+class _StubDataService:
+    """Minimal duck-typed stand-in for DataService -- _process_sample_row
+    only touches self._ctx / self._get_dataset / self._is_metadata_only_request
+    / self._is_nan_value, so a full instance with real gRPC/context wiring
+    isn't needed to exercise its label/pred string-handling branches."""
+
+    def __init__(self):
+        self._ctx = None
+
+    def _get_dataset(self, origin):
+        return None
+
+    def _is_metadata_only_request(self, request):
+        return DataService._is_metadata_only_request(self, request)
+
+    def _is_nan_value(self, value):
+        return DataService._is_nan_value(self, value)
+
+
+def _find_stat(data_stats, name):
+    return next((s for s in data_stats if s.name == name), None)
+
+
+class TestProcessSampleRowTextLabelPrediction(unittest.TestCase):
+    def setUp(self):
+        self.service = _StubDataService()
+        # resize_width/height > 0 so _is_metadata_only_request returns False
+        # and the label/prediction branches actually run (an all-defaults
+        # request is treated as a metadata-only histogram sweep that skips
+        # them entirely).
+        self.request = pb2.DataSamplesRequest(resize_width=32, resize_height=32)
+
+    def _process(self, row):
+        data_record = DataService._process_sample_row(self.service, (row, self.request, None))
+        self.assertIsNotNone(data_record, "record processing raised internally (see logged exception)")
+        return list(data_record.data_stats)
+
+    def test_text_label_and_prediction_become_string_data_stats(self):
+        row = {
+            SampleStatsEx.SAMPLE_ID.value: "1",
+            SampleStatsEx.ORIGIN.value: "train_loader",
+            SampleStatsEx.TASK_TYPE.value: "classification",
+            SampleStatsEx.TARGET.value: "The reference/target answer.",
+            SampleStatsEx.PREDICTION.value: "The model's generated reply.",
+        }
+        data_stats = self._process(row)
+
+        target_stat = _find_stat(data_stats, "target")
+        pred_stat = _find_stat(data_stats, "pred")
+
+        self.assertIsNotNone(target_stat)
+        self.assertEqual(target_stat.type, "string")
+        self.assertEqual(target_stat.value_string, "The reference/target answer.")
+
+        self.assertIsNotNone(pred_stat)
+        self.assertEqual(pred_stat.type, "string")
+        self.assertEqual(pred_stat.value_string, "The model's generated reply.")
+
+    def test_multi_sentence_text_without_trailing_period_is_not_dropped(self):
+        # Regression: the old is_label_empty heuristic would misclassify
+        # this as a "file path" (contains a '.', doesn't end in one) and
+        # treat it as empty, discarding it instead of emitting a stat.
+        row = {
+            SampleStatsEx.SAMPLE_ID.value: "2",
+            SampleStatsEx.ORIGIN.value: "train_loader",
+            SampleStatsEx.TASK_TYPE.value: "classification",
+            SampleStatsEx.TARGET.value: "Sure. Here is a multi-sentence answer without a period at the end",
+            SampleStatsEx.PREDICTION.value: "Sure. Here is the reply without a period at the end",
+        }
+        data_stats = self._process(row)
+
+        target_stat = _find_stat(data_stats, "target")
+        pred_stat = _find_stat(data_stats, "pred")
+        self.assertIsNotNone(target_stat)
+        self.assertEqual(target_stat.type, "string")
+        self.assertIsNotNone(pred_stat)
+        self.assertEqual(pred_stat.type, "string")
+
+    def test_numeric_classification_label_and_prediction_unaffected(self):
+        """Backward-compat guard: existing numeric classification labels
+        must still produce scalar DataStats, unaffected by the text path."""
+        row = {
+            SampleStatsEx.SAMPLE_ID.value: "3",
+            SampleStatsEx.ORIGIN.value: "train_loader",
+            SampleStatsEx.TASK_TYPE.value: "classification",
+            SampleStatsEx.TARGET.value: 1,
+            SampleStatsEx.PREDICTION.value: 0,
+        }
+        data_stats = self._process(row)
+
+        target_stat = _find_stat(data_stats, "target")
+        pred_stat = _find_stat(data_stats, "pred")
+        self.assertIsNotNone(target_stat)
+        self.assertEqual(target_stat.type, "scalar")
+        self.assertEqual(list(target_stat.value), [1.0])
+        self.assertIsNotNone(pred_stat)
+        self.assertEqual(pred_stat.type, "scalar")
+        self.assertEqual(list(pred_stat.value), [0.0])
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/trainer/test_trainer_tools.py

@@ -12,6 +12,7 @@
     _class_ids,
     _get_input_tensor_for_sample,
     _labels_from_mask_path_histogram,
+    _numeric_or_text_list,
     execute_df_operation,
     force_kill_all_python_processes,
     generate_overview,
@@ -110,6 +111,8 @@ def __init__(self, **kwargs):
                 self.task_type = kwargs.get("task_type", "")
                 self.sample_label = []
                 self.sample_prediction = []
+                self.sample_label_text = []
+                self.sample_prediction_text = []
 
         class _FakeSampleStats:
             def __init__(self):
@@ -143,6 +146,54 @@ def __init__(self):
             self.assertEqual(list(seg_stats.records[0].sample_label), [0, 1, 2])
             self.assertEqual(list(seg_stats.records[0].sample_prediction), [1, 2, 4])
 
+    def test_numeric_or_text_list(self):
+        self.assertEqual(_numeric_or_text_list("a generated reply"), ([], ["a generated reply"]))
+        self.assertEqual(_numeric_or_text_list(3), ([3], []))
+        self.assertEqual(_numeric_or_text_list([3]), ([3], []))
+        self.assertEqual(_numeric_or_text_list(np.array([3])), ([3], []))
+        self.assertEqual(_numeric_or_text_list(["only one string in a list"]), ([], ["only one string in a list"]))
+
+    def test_get_data_set_representation_text_label_and_prediction(self):
+        """A generative task (e.g. an LLM's reply as prediction, a reference
+        string as label) must route through the new text fields instead of
+        crashing on int(target)/int(pred)."""
+        exp = _Experiment()
+
+        class _FakeRecord:
+            def __init__(self, **kwargs):
+                self.sample_id = kwargs.get("sample_id")
+                self.sample_last_loss = kwargs.get("sample_last_loss", -1.0)
+                self.sample_discarded = kwargs.get("sample_discarded", False)
+                self.task_type = kwargs.get("task_type", "")
+                self.sample_label = []
+                self.sample_prediction = []
+                self.sample_label_text = []
+                self.sample_prediction_text = []
+
+        class _FakeSampleStats:
+            def __init__(self):
+                self.sample_count = 0
+                self.task_type = ""
+                self.records = []
+
+        text_rows = [
+            {
+                "sample_id": "30",
+                "prediction_loss": 0.1,
+                "label": "The reference/target answer.",
+                "prediction_raw": "The model's generated reply.",
+                "discarded": False,
+            },
+        ]
+
+        with patch("weightslab.trainer.trainer_tools.pb2.SampleStatistics", side_effect=_FakeSampleStats), \
+             patch("weightslab.trainer.trainer_tools.pb2.RecordMetadata", side_effect=lambda **kw: _FakeRecord(**kw)):
+            text_stats = get_data_set_representation(_SimpleDataset(text_rows), exp)
+            self.assertEqual(list(text_stats.records[0].sample_label), [])
+            self.assertEqual(list(text_stats.records[0].sample_prediction), [])
+            self.assertEqual(list(text_stats.records[0].sample_label_text), ["The reference/target answer."])
+            self.assertEqual(list(text_stats.records[0].sample_prediction_text), ["The model's generated reply."])
+
     def test_load_raw_image_from_images_and_tensor_input(self):
         with tempfile.TemporaryDirectory() as tmp:
             p = os.path.join(tmp, "img.png")
@@ -167,6 +218,34 @@ def test_get_input_tensor_for_sample(self):
         tensor = _get_input_tensor_for_sample(_IndexDataset(), sample_id=0, device="cpu")
         self.assertEqual(tuple(tensor.shape), (1, 3, 8, 8))
 
+    def test_process_sample_text_only_dataset_returns_empty_without_error(self):
+        """A generative/RLHF prompt dataset's items are text, not images --
+        process_sample (used by the GetSamples image-preview endpoint) must
+        return an empty/placeholder result instead of crashing on
+        torch.tensor(a_string) (confirmed empirically: raises 'new(): invalid
+        data type str')."""
+
+        class _TextDataset:
+            task_type = "generation"
+
+            def _getitem_raw(self, id):
+                return "flip a coin", id, "Certainly! Let's flip a fair coin..."
+
+        sid, transformed, raw, cls_label, mask_bytes, pred_bytes = process_sample(
+            sid=0,
+            dataset=_TextDataset(),
+            do_resize=False,
+            resize_dims=(8, 8),
+            experiment=_Experiment(),
+        )
+
+        self.assertEqual(sid, 0)
+        self.assertIsNone(transformed)
+        self.assertIsNone(raw)
+        self.assertEqual(cls_label, -1)
+        self.assertEqual(mask_bytes, b"")
+        self.assertEqual(pred_bytes, b"")
+
     def test_process_sample_classification_and_force_kill(self):
         exp = _Experiment()
         ds = _RawDataset()

weightslab/proto/experiment_service.proto

@@ -278,6 +278,13 @@ message RecordMetadata {
   repeated TaskField extra_fields = 7;
   bytes prediction_raw = 9;
   string task_type = 10;
+  // Text label/prediction (e.g. a generative model's target/reference text
+  // and its generated reply) -- additive, alongside the int32 fields above
+  // rather than replacing them, so existing classification/segmentation/
+  // detection consumers that populate sample_label/sample_prediction are
+  // completely unaffected. Empty/unset for numeric-label tasks.
+  repeated string sample_label_text = 11;
+  repeated string sample_prediction_text = 12;
 }
 
 message SampleStatistics {