feat: implement cost sensitive labels (#55)

Author: jackgerrits

src/main.cpp

@@ -4,6 +4,7 @@
 #include "vw/core/cache.h"
 #include "vw/core/cb.h"
 #include "vw/core/constant.h"
+#include "vw/core/cost_sensitive.h"
 #include "vw/core/example.h"
 #include "vw/core/global_data.h"
 #include "vw/core/guard.h"
@@ -645,6 +646,15 @@ struct py_simple_label
   float initial;
 };
 
+bool is_shared(const VW::cs_label& label)
+{
+  const auto& costs = label.costs;
+  if (costs.size() != 1) { return false; }
+  if (costs[0].class_index != 0) { return false; }
+  if (costs[0].x != -FLT_MAX) { return false; }
+  return true;
+}
+
 }  // namespace
 
 PYBIND11_MODULE(_core, m)
@@ -790,6 +800,9 @@ PYBIND11_MODULE(_core, m)
           py::kw_only(), py::arg("label") = py::none(), py::arg("weight") = 1.f, py::arg("shared") = false, R"docstring(
     A label representing a contextual bandit problem.
 
+    .. note::
+      Currently the label can only contain 1 or 0 cb costs. There is a mode in vw for CB (non-adf) that allows for multiple cb_classes per example, but it is not currently accessible via direct label access. If creating examples/labels from parsed input it should still work as expected. If you need this feature, please open an issue on the github repo.
+
     Args:
       label (Optional[Union[Tuple[float, float], Tuple[int, float, float]]): This is (action, cost, probability). The same rules as VW apply for if the action can be left out of the tuple.
       weight (float): The weight of the example.
@@ -841,6 +854,63 @@ PYBIND11_MODULE(_core, m)
     The weight of the example.
 )docstring");
 
+  py::class_<VW::cs_label>(m, "CSLabel")
+      .def(py::init(
+               [](std::optional<std::vector<std::tuple<float, float>>> costs, bool is_shared)
+               {
+                 auto label = std::make_unique<VW::cs_label>();
+                 if (is_shared)
+                 {
+                   if (costs.has_value())
+                   {
+                     throw std::invalid_argument("Shared examples cannot have action, cost, or probability.");
+                   }
+
+                   label->costs.emplace_back(-FLT_MAX, 0, 0.f, 0.f);
+                   return label;
+                 }
+
+                 if (costs.has_value())
+                 {
+                   for (auto& [class_index, cost] : *costs) { label->costs.emplace_back(cost, class_index, 0.f, 0.f); }
+                 }
+
+                 return label;
+               }),
+          py::kw_only(), py::arg("costs") = py::none(), py::arg("shared") = false, R"docstring(
+    A label representing a cost sensitive classification problem.
+
+    Args:
+      costs (Optional[List[Tuple[int, float]]]): List of classes and costs. If there is no label, this should be None.
+      shared (bool): Whether the example represents the shared context
+)docstring")
+      .def_property_readonly(
+          "shared", [](VW::cs_label& l) -> bool { return is_shared(l); },
+          R"docstring(
+    Whether the example represents the shared context.
+)docstring")
+      .def_property(
+          "costs",
+          [](VW::cs_label& l) -> std::optional<std::vector<std::tuple<uint32_t, float>>>
+          {
+            if (is_shared(l)) { return std::nullopt; }
+
+            std::vector<std::tuple<uint32_t, float>> costs;
+            costs.reserve(l.costs.size());
+            for (auto& cost : l.costs) { costs.emplace_back(cost.class_index, cost.x); }
+            return costs;
+          },
+          [](VW::cs_label& l, const std::vector<std::tuple<uint32_t, float>>& label)
+          {
+            if (is_shared(l)) { throw std::invalid_argument("Shared examples cannot have costs."); }
+
+            l.costs.clear();
+            for (auto& [class_index, cost] : label) { l.costs.emplace_back(cost, class_index, 0.f, 0.f); }
+          },
+          R"docstring(
+    The costs for the example. The format of the costs is (class_index, cost).
+)docstring");
+
   py::class_<VW::example, std::shared_ptr<VW::example>>(m, "Example")
       .def(py::init(
           []()
@@ -851,7 +921,7 @@ PYBIND11_MODULE(_core, m)
       .def("_is_newline", [](VW::example& ex) -> bool { return ex.is_newline; })
       .def("_get_label",
           [](VW::example& ex, VW::label_type_t label_type)
-              -> std::variant<py_simple_label, VW::multiclass_label, VW::cb_label, std::monostate>
+              -> std::variant<py_simple_label, VW::multiclass_label, VW::cb_label, VW::cs_label, std::monostate>
           {
             switch (label_type)
             {
@@ -864,6 +934,8 @@ PYBIND11_MODULE(_core, m)
                 return ex.l.multi;
               case VW::label_type_t::CB:
                 return ex.l.cb;
+              case VW::label_type_t::CS:
+                return ex.l.cs;
               case VW::label_type_t::NOLABEL:
                 return std::monostate();
               default:
@@ -872,7 +944,7 @@ PYBIND11_MODULE(_core, m)
           })
       .def("_set_label",
           [](VW::example& ex,
-              const std::variant<py_simple_label*, VW::multiclass_label*, VW::cb_label*, std::monostate>& label) -> void
+              const std::variant<py_simple_label*, VW::multiclass_label*, VW::cb_label*, VW::cs_label*, std::monostate>& label) -> void
           {
             std::visit(
                 overloaded{
@@ -887,6 +959,7 @@ PYBIND11_MODULE(_core, m)
                     },
                     [&](VW::multiclass_label* multiclass_label) { ex.l.multi = *multiclass_label; },
                     [&](VW::cb_label* cb_label) { ex.l.cb = *cb_label; },
+                    [&](VW::cs_label* cs_label) { ex.l.cs = *cs_label; },
                 },
                 label);
           });

src/vowpal_wabbit_next/__init__.py

@@ -7,7 +7,7 @@
 from .delta import ModelDelta, calculate_delta, apply_delta, merge_deltas
 from .cli_driver import CLIError, run_cli_driver
 from .prediction_type import PredictionType
-from .labels import LabelType, SimpleLabel, MulticlassLabel, CBLabel
+from .labels import LabelType, SimpleLabel, MulticlassLabel, CBLabel, CSLabel
 
 
 VW_COMMIT: str = _vw_commit
@@ -38,4 +38,5 @@
     "SimpleLabel",
     "MulticlassLabel",
     "CBLabel",
+    "CSLabel",
 ]

src/vowpal_wabbit_next/_core.pyi

@@ -4,6 +4,7 @@ import typing
 
 __all__ = [
     "CBLabel",
+    "CSLabel",
     "DenseParameters",
     "Example",
     "LabelType",
@@ -20,6 +21,9 @@ class CBLabel():
         """
         A label representing a contextual bandit problem.
 
+        .. note::
+          Currently the label can only contain 1 or 0 cb costs. There is a mode in vw for CB (non-adf) that allows for multiple cb_classes per example, but it is not currently accessible via direct label access. If creating examples/labels from parsed input it should still work as expected. If you need this feature, please open an issue on the github repo.
+
         Args:
           label (Optional[Union[Tuple[float, float], Tuple[int, float, float]]): This is (action, cost, probability). The same rules as VW apply for if the action can be left out of the tuple.
           weight (float): The weight of the example.
@@ -57,13 +61,42 @@ class CBLabel():
         The weight of the example.
         """
     pass
+class CSLabel():
+    def __init__(self, *, costs: typing.Optional[typing.List[typing.Tuple[float, float]]] = None, shared: bool = False) -> None: 
+        """
+        A label representing a cost sensitive classification problem.
+
+        Args:
+          costs (Optional[List[Tuple[int, float]]]): List of classes and costs. If there is no label, this should be None.
+          shared (bool): Whether the example represents the shared context
+        """
+    @property
+    def costs(self) -> typing.Optional[typing.List[typing.Tuple[int, float]]]:
+        """
+            The costs for the example. The format of the costs is (class_index, cost).
+
+        :type: typing.Optional[typing.List[typing.Tuple[int, float]]]
+        """
+    @costs.setter
+    def costs(self, arg1: typing.List[typing.Tuple[int, float]]) -> None:
+        """
+        The costs for the example. The format of the costs is (class_index, cost).
+        """
+    @property
+    def shared(self) -> bool:
+        """
+            Whether the example represents the shared context.
+
+        :type: bool
+        """
+    pass
 class DenseParameters():
     pass
 class Example():
     def __init__(self) -> None: ...
-    def _get_label(self, arg0: LabelType) -> typing.Union[SimpleLabel, MulticlassLabel, CBLabel, None]: ...
+    def _get_label(self, arg0: LabelType) -> typing.Union[SimpleLabel, MulticlassLabel, CBLabel, CSLabel, None]: ...
     def _is_newline(self) -> bool: ...
-    def _set_label(self, arg0: typing.Union[SimpleLabel, MulticlassLabel, CBLabel, None]) -> None: ...
+    def _set_label(self, arg0: typing.Union[SimpleLabel, MulticlassLabel, CBLabel, CSLabel, None]) -> None: ...
     pass
 class LabelType():
     def __eq__(self, other: object) -> bool: ...

src/vowpal_wabbit_next/example.py

@@ -1,5 +1,11 @@
 from vowpal_wabbit_next import _core
-from vowpal_wabbit_next.labels import LabelType, SimpleLabel, MulticlassLabel, CBLabel
+from vowpal_wabbit_next.labels import (
+    LabelType,
+    SimpleLabel,
+    MulticlassLabel,
+    CBLabel,
+    CSLabel,
+)
 from typing import Optional, Union
 
 
@@ -27,7 +33,7 @@ def __init__(
             self._example = _core.Example()
             self.label_type = LabelType.NoLabel
 
-    def get_label(self) -> Union[SimpleLabel, MulticlassLabel, CBLabel, None]:
+    def get_label(self) -> Union[SimpleLabel, MulticlassLabel, CBLabel, CSLabel, None]:
         """Get the label of the example.
 
         Returns:
@@ -36,7 +42,7 @@ def get_label(self) -> Union[SimpleLabel, MulticlassLabel, CBLabel, None]:
         return self._example._get_label(self.label_type)
 
     def set_label(
-        self, label: Union[SimpleLabel, MulticlassLabel, CBLabel, None]
+        self, label: Union[SimpleLabel, MulticlassLabel, CBLabel, CSLabel, None]
     ) -> None:
         """Set the label of the example.
 
@@ -52,6 +58,8 @@ def set_label(
             self.label_type = LabelType.Multiclass
         elif isinstance(label, CBLabel):
             self.label_type = LabelType.CB
+        elif isinstance(label, CSLabel):
+            self.label_type = LabelType.CS
         elif label is None:
             self.label_type = LabelType.NoLabel
         else:

src/vowpal_wabbit_next/labels.py

@@ -4,3 +4,4 @@
 SimpleLabel = _core.SimpleLabel
 MulticlassLabel = _core.MulticlassLabel
 CBLabel = _core.CBLabel
+CSLabel = _core.CSLabel

tests/test_labels.py

@@ -43,3 +43,39 @@ def test_cb_label() -> None:
 
     example.set_label(None)
     assert example.get_label() is None
+
+
+def test_cs_label() -> None:
+    example = vw.Example()
+    assert example.get_label() is None
+    example.set_label(vw.CSLabel(costs=[(1, 1.2)]))
+    assert isinstance(example.get_label(), vw.CSLabel)
+    assert len(example.get_label().costs) == 1
+    assert example.get_label().costs[0] == pytest.approx((1, 1.2))
+
+    with pytest.raises(ValueError):
+        example.set_label(vw.CSLabel(costs=[(1, 1.2)], shared=True))
+
+    example.set_label(vw.CSLabel(shared=True))
+    assert isinstance(example.get_label(), vw.CSLabel)
+    assert example.get_label().costs is None
+    assert example.get_label().shared == True
+
+    example.set_label(None)
+    assert example.get_label() is None
+
+
+def test_cs_label_parsed() -> None:
+    workspace = vw.Workspace(["--csoaa_ldf=mc"])
+    parser = vw.TextFormatParser(workspace)
+    shared_ex = parser.parse_line("shared | a b c")
+    labeled_ex = parser.parse_line("1:0.5 | a b c")
+    assert isinstance(shared_ex.get_label(), vw.CSLabel)
+    assert isinstance(labeled_ex.get_label(), vw.CSLabel)
+
+    assert shared_ex.get_label().costs is None
+    assert shared_ex.get_label().shared == True
+
+    assert len(labeled_ex.get_label().costs) == 1
+    assert labeled_ex.get_label().costs[0] == pytest.approx((1, 0.5))
+    assert labeled_ex.get_label().shared == False