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)