No vw-related logic in PickBestFeaturizer

src/learn_to_pick/__init__.py

@@ -53,5 +53,4 @@ def configure_logger() -> None:
     "VwPolicy",
     "VwLogger",
     "embed",
-    "stringify_embedding",
 ]

src/learn_to_pick/base.py

@@ -13,11 +13,13 @@
     Type,
     TypeVar,
     Union,
+    Callable,
 )
 
 from learn_to_pick.metrics import MetricsTrackerAverage, MetricsTrackerRollingWindow
 from learn_to_pick.model_repository import ModelRepository
 from learn_to_pick.vw_logger import VwLogger
+from learn_to_pick.features import Featurized, DenseFeatures, SparseFeatures
 
 if TYPE_CHECKING:
     import vowpal_wabbit_next as vw
@@ -87,10 +89,6 @@ def EmbedAndKeep(anything: Any) -> Any:
 # helper functions
 
 
-def _stringify_embedding(embedding: List) -> str:
-    return " ".join([f"{i}:{e}" for i, e in enumerate(embedding)])
-
-
 def _parse_lines(parser: "vw.TextFormatParser", input_str: str) -> List["vw.Example"]:
     return [parser.parse_line(line) for line in input_str.split("\n")]
 
@@ -108,7 +106,7 @@ def get_based_on_and_to_select_from(inputs: Dict[str, Any]) -> Tuple[Dict, Dict]
         )
 
     based_on = {
-        k: inputs[k].value if isinstance(inputs[k].value, list) else [inputs[k].value]
+        k: inputs[k].value if isinstance(inputs[k].value, list) else inputs[k].value
         for k in inputs.keys()
         if isinstance(inputs[k], _BasedOn)
     }
@@ -165,36 +163,38 @@ def __init__(
         model_repo: ModelRepository,
         vw_cmd: List[str],
         featurizer: Featurizer,
+        formatter: Callable,
         vw_logger: VwLogger,
-        *args: Any,
         **kwargs: Any,
     ):
-        super().__init__(*args, **kwargs)
+        super().__init__(**kwargs)
         self.model_repo = model_repo
         self.vw_cmd = vw_cmd
         self.workspace = self.model_repo.load(vw_cmd)
         self.featurizer = featurizer
+        self.formatter = formatter
         self.vw_logger = vw_logger
 
+    def format(self, event):
+        return self.formatter(*self.featurizer.featurize(event))
+
     def predict(self, event: TEvent) -> Any:
         import vowpal_wabbit_next as vw
 
         text_parser = vw.TextFormatParser(self.workspace)
-        return self.workspace.predict_one(
-            _parse_lines(text_parser, self.featurizer.format(event))
-        )
+        return self.workspace.predict_one(_parse_lines(text_parser, self.format(event)))
 
     def learn(self, event: TEvent) -> None:
         import vowpal_wabbit_next as vw
 
-        vw_ex = self.featurizer.format(event)
+        vw_ex = self.format(event)
         text_parser = vw.TextFormatParser(self.workspace)
         multi_ex = _parse_lines(text_parser, vw_ex)
         self.workspace.learn_one(multi_ex)
 
     def log(self, event: TEvent) -> None:
         if self.vw_logger.logging_enabled():
-            vw_ex = self.featurizer.format(event)
+            vw_ex = self.format(event)
             self.vw_logger.log(vw_ex)
 
     def save(self) -> None:
@@ -206,7 +206,7 @@ def __init__(self, *args: Any, **kwargs: Any):
         pass
 
     @abstractmethod
-    def format(self, event: TEvent) -> Any:
+    def featurize(self, event: TEvent) -> Any:
         ...
 
 
@@ -486,70 +486,59 @@ def run(self, *args, **kwargs) -> Dict[str, Any]:
 
 
 def _embed_string_type(
-    item: Union[str, _Embed], model: Any, namespace: Optional[str] = None
-) -> Dict[str, Union[str, List[str]]]:
+    item: Union[str, _Embed], model: Any, namespace: str
+) -> Featurized:
     """Helper function to embed a string or an _Embed object."""
     import re
 
-    keep_str = ""
+    result = Featurized()
     if isinstance(item, _Embed):
-        encoded = _stringify_embedding(model.encode(item.value))
-        # TODO these should be moved to pick_best
+        result[namespace] = DenseFeatures(model.encode(item.value))
         if item.keep:
-            keep_str = item.value.replace(" ", "_") + " "
-            keep_str = re.sub(r"[\t\n\r\f\v]+", " ", keep_str)
+            keep_str = item.value.replace(" ", "_")
+            result[namespace] = {"default_ft": re.sub(r"[\t\n\r\f\v]+", " ", keep_str)}
     elif isinstance(item, str):
         encoded = item.replace(" ", "_")
-        encoded = re.sub(r"[\t\n\r\f\v]+", " ", encoded)
+        result[namespace] = {"default_ft": re.sub(r"[\t\n\r\f\v]+", " ", encoded)}
     else:
         raise ValueError(f"Unsupported type {type(item)} for embedding")
 
-    if namespace is None:
-        raise ValueError(
-            "The default namespace must be provided when embedding a string or _Embed object."
-        )
-
-    return {namespace: keep_str + encoded}
+    return result
 
 
-def _embed_dict_type(item: Dict, model: Any) -> Dict[str, Any]:
+def _embed_dict_type(item: Dict, model: Any) -> Featurized:
     """Helper function to embed a dictionary item."""
-    inner_dict: Dict = {}
+    result = Featurized()
     for ns, embed_item in item.items():
         if isinstance(embed_item, list):
-            inner_dict[ns] = []
-            for embed_list_item in embed_item:
-                embedded = _embed_string_type(embed_list_item, model, ns)
-                inner_dict[ns].append(embedded[ns])
+            for idx, embed_list_item in enumerate(embed_item):
+                result.merge(_embed_string_type(embed_list_item, model, f"{ns}_{idx}"))
         else:
-            inner_dict.update(_embed_string_type(embed_item, model, ns))
-    return inner_dict
+            result.merge(_embed_string_type(embed_item, model, ns))
+    return result
 
 
 def _embed_list_type(
     item: list, model: Any, namespace: Optional[str] = None
-) -> List[Dict[str, Union[str, List[str]]]]:
-    ret_list: List = []
+) -> List[Featurized]:
+    result = []
     for embed_item in item:
         if isinstance(embed_item, dict):
-            ret_list.append(_embed_dict_type(embed_item, model))
+            result.append(_embed_dict_type(embed_item, model))
         elif isinstance(embed_item, list):
-            item_embedding = _embed_list_type(embed_item, model, namespace)
-            # Get the first key from the first dictionary
-            first_key = next(iter(item_embedding[0]))
-            # Group the values under that key
-            grouping = {first_key: [item[first_key] for item in item_embedding]}
-            ret_list.append(grouping)
+            result.append(Featurized())
+            for idx, embed_list_item in enumerate(embed_item):
+                result[-1].merge(_embed_string_type(embed_list_item, model, f"{idx}"))
         else:
-            ret_list.append(_embed_string_type(embed_item, model, namespace))
-    return ret_list
+            result.append(_embed_string_type(embed_item, model, namespace))
+    return result
 
 
 def embed(
     to_embed: Union[Union[str, _Embed], Dict, List[Union[str, _Embed]], List[Dict]],
     model: Any,
     namespace: Optional[str] = None,
-) -> List[Dict[str, Union[str, List[str]]]]:
+) -> Union[Featurized, List[Featurized]]:
     """
     Embeds the actions or context using the SentenceTransformer model (or a model that has an `encode` function)
 
@@ -563,9 +552,9 @@ def embed(
     if (isinstance(to_embed, _Embed) and isinstance(to_embed.value, str)) or isinstance(
         to_embed, str
     ):
-        return [_embed_string_type(to_embed, model, namespace)]
+        return _embed_string_type(to_embed, model, namespace)
     elif isinstance(to_embed, dict):
-        return [_embed_dict_type(to_embed, model)]
+        return _embed_dict_type(to_embed, model)
     elif isinstance(to_embed, list):
         return _embed_list_type(to_embed, model, namespace)
     else:

src/learn_to_pick/features.py

@@ -0,0 +1,36 @@
+from typing import Union, Optional, Dict, List
+import numpy as np
+
+
+class SparseFeatures(dict):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+
+class DenseFeatures(list):
+    def __init__(self, *args, **kwargs):
+        super().__init__(np.array(*args, **kwargs))
+
+
+class Featurized:
+    def __init__(
+        self,
+        sparse: Optional[Dict[str, SparseFeatures]] = None,
+        dense: Optional[Dict[str, DenseFeatures]] = None,
+    ):
+        self.sparse = sparse or {}
+        self.dense = dense or {}
+
+    def __setitem__(self, key, value):
+        if isinstance(value, Dict):
+            self.sparse[key] = SparseFeatures(value)
+        elif isinstance(value, List) or isinstance(value, np.ndarray):
+            self.dense[key] = DenseFeatures(value)
+        else:
+            raise ValueError(
+                f"Cannot convert {type(value)} to either DenseFeatures or SparseFeatures"
+            )
+
+    def merge(self, other):
+        self.sparse.update(other.sparse)
+        self.dense.update(other.dense)