Merge pull request #31 from VowpalWabbit/byom

Add pytorch policy

README.md

@@ -29,7 +29,7 @@ Note: all code examples presented here can be found in `notebooks/readme.ipynb`
 - Use a custom score function to grade the decision.
 - Directly specify the score manually and asynchronously.
 
-The beauty of `learn_to_pick` is its flexibility. Whether you're a fan of VowpalWabbit or prefer PyTorch (coming soon), the library can seamlessly integrate with both, allowing them to be the brain behind your decisions.
+The beauty of `learn_to_pick` is its flexibility. Whether you're a fan of VowpalWabbit or prefer PyTorch, the library can seamlessly integrate with both, allowing them to be the brain behind your decisions.
 
 ## Installation
 
@@ -43,6 +43,8 @@ The `PickBest` scenario should be used when:
 - Only one option is optimal for a specific criteria or context
 - There exists a mechanism to provide feedback on the suitability of the chosen option for the specific criteria
 
+### Scorer
+
 Example usage with llm default scorer:
 
 ```python
@@ -113,7 +115,46 @@ dummy_score = 1
 picker.update_with_delayed_score(dummy_score, result)
 ```
 
-`PickBest` is highly configurable to work with a VowpalWabbit decision making policy, a PyTorch decision making policy (coming soon), or with a custom user defined decision making policy
+### Using Pytorch policy
+
+Example usage with a Pytorch policy:
+```python
+from learn_to_pick import PyTorchPolicy
+
+pytorch_picker = learn_to_pick.PickBest.create(
+    policy=PyTorchPolicy(), selection_scorer=CustomSelectionScorer())
+
+pytorch_picker.run(
+    pick = learn_to_pick.ToSelectFrom(["option1", "option2"]),
+    criteria = learn_to_pick.BasedOn("some criteria")
+)
+```
+
+Example usage with a custom Pytorch policy:
+You can alway create a custom Pytorch policy by implementing the Policy interface
+
+```python
+class CustomPytorchPolicy(Policy):
+    def __init__(self, **kwargs: Any):
+        ...
+
+    def predict(self, event: TEvent) -> Any:
+        ...
+
+    def learn(self, event: TEvent) -> None:
+        ...
+
+    def log(self, event: TEvent) -> None:
+        ...
+
+    def save(self) -> None:
+        ...
+
+pytorch_picker = learn_to_pick.PickBest.create(
+    policy=CustomPytorchPolicy(), selection_scorer=CustomSelectionScorer())
+```
+
+`PickBest` is highly configurable to work with a VowpalWabbit decision making policy, a PyTorch decision making policy, or with a custom user defined decision making policy
 
 The main thing that needs to be decided from the get-go is:
 
@@ -134,7 +175,8 @@ In all three cases, when a score is calculated or provided, the decision making
 ## Example Notebooks
 
 - `readme.ipynb` showcases all examples shown in this README
-- `news_recommendation.ipynb` showcases a personalization scenario where we have to pick articles for specific users
+- `news_recommendation.ipynb` showcases a personalization scenario where we have to pick articles for specific users with VowpalWabbit policy
+- `news_recommendation_pytorch.ipynb` showcases the same personalization scenario where we have to pick articles for specific users with Pytorch policy
 - `prompt_variable_injection.ipynb` showcases learned prompt variable injection and registering callback functionality
 
 ### Advanced Usage
@@ -183,7 +225,7 @@ class CustomSelectionScorer(learn_to_pick.SelectionScorer):
         # inputs: the inputs to the picker in Dict[str, Any] format
         # picked: the selection that was made by the policy
         # event: metadata that can be used to determine the score if needed
-        
+
         # scoring logic goes here
 
         dummy_score = 1.0

setup.py

@@ -1,5 +1,4 @@
 from setuptools import setup, find_packages
-import os
 
 with open("README.md", "r", encoding="UTF-8") as fh:
     long_description = fh.read()

src/learn_to_pick/__init__.py

@@ -5,12 +5,9 @@
     BasedOn,
     Embed,
     Featurizer,
-    ModelRepository,
     Policy,
     SelectionScorer,
     ToSelectFrom,
-    VwPolicy,
-    VwLogger,
     embed,
 )
 from learn_to_pick.pick_best import (
@@ -22,6 +19,14 @@
 )
 
 
+from learn_to_pick.vw.policy import VwPolicy
+from learn_to_pick.vw.model_repository import ModelRepository
+from learn_to_pick.vw.logger import VwLogger
+
+from learn_to_pick.pytorch.policy import PyTorchPolicy
+from learn_to_pick.pytorch.feature_embedder import PyTorchFeatureEmbedder
+
+
 def configure_logger() -> None:
     logger = logging.getLogger(__name__)
     logger.setLevel(logging.INFO)
@@ -48,9 +53,11 @@ def configure_logger() -> None:
     "SelectionScorer",
     "AutoSelectionScorer",
     "Featurizer",
-    "ModelRepository",
     "Policy",
+    "PyTorchPolicy",
+    "PyTorchFeatureEmbedder",
+    "embed",
+    "ModelRepository",
     "VwPolicy",
     "VwLogger",
-    "embed",
 ]

src/learn_to_pick/base.py

@@ -10,15 +10,12 @@
     List,
     Optional,
     Tuple,
-    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
 from enum import Enum
 
@@ -89,10 +86,6 @@ def EmbedAndKeep(anything: Any) -> Any:
 # helper functions
 
 
-def _parse_lines(parser: "vw.TextFormatParser", input_str: str) -> List["vw.Example"]:
-    return [parser.parse_line(line) for line in input_str.split("\n")]
-
-
 def filter_inputs(inputs: Dict[str, Any], role: Role) -> Dict[str, Any]:
     return {
         k: v.value
@@ -144,50 +137,6 @@ def save(self) -> None:
         pass
 
 
-class VwPolicy(Policy):
-    def __init__(
-        self,
-        model_repo: ModelRepository,
-        vw_cmd: List[str],
-        featurizer: Featurizer,
-        formatter: Callable,
-        vw_logger: VwLogger,
-        **kwargs: Any,
-    ):
-        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.format(event)))
-
-    def learn(self, event: TEvent) -> None:
-        import vowpal_wabbit_next as vw
-
-        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.format(event)
-            self.vw_logger.log(vw_ex)
-
-    def save(self) -> None:
-        self.model_repo.save(self.workspace)
-
-
 class Featurizer(Generic[TEvent], ABC):
     def __init__(self, *args: Any, **kwargs: Any):
         pass

src/learn_to_pick/pick_best.py

@@ -7,6 +7,10 @@
 import numpy as np
 
 from learn_to_pick import base
+from learn_to_pick.vw.policy import VwPolicy
+from learn_to_pick.vw.model_repository import ModelRepository
+from learn_to_pick.vw.logger import VwLogger
+
 
 logger = logging.getLogger(__name__)
 
@@ -333,14 +337,14 @@ def create_policy(
 
         vw_cmd = interactions + vw_cmd
 
-        return base.VwPolicy(
-            model_repo=base.ModelRepository(
+        return VwPolicy(
+            model_repo=ModelRepository(
                 model_save_dir, with_history=True, reset=reset_model
             ),
             vw_cmd=vw_cmd,
             featurizer=featurizer,
             formatter=formatter,
-            vw_logger=base.VwLogger(rl_logs),
+            vw_logger=VwLogger(rl_logs),
         )
 
     def _default_policy(self):

src/learn_to_pick/pytorch/feature_embedder.py

@@ -0,0 +1,69 @@
+from sentence_transformers import SentenceTransformer
+import torch
+from torch import Tensor
+
+from learn_to_pick import PickBestFeaturizer
+from learn_to_pick.base import Event
+from learn_to_pick.features import SparseFeatures
+from typing import Any, Tuple, TypeVar, Union
+
+TEvent = TypeVar("TEvent", bound=Event)
+
+
+class PyTorchFeatureEmbedder:
+    def __init__(self, model: Any = None):
+        if model is None:
+            model = SentenceTransformer("all-MiniLM-L6-v2")
+
+        self.model = model
+        self.featurizer = PickBestFeaturizer(auto_embed=False)
+
+    def encode(self, to_encode: str) -> Tensor:
+        embeddings = self.model.encode(to_encode, convert_to_tensor=True)
+        normalized = torch.nn.functional.normalize(embeddings)
+        return normalized
+
+    def convert_features_to_text(self, sparse_features: SparseFeatures) -> str:
+        results = []
+        for ns, obj in sparse_features.items():
+            value = obj.get("default_ft", "")
+            results.append(f"{ns}={value}")
+        return " ".join(results)
+
+    def format(
+        self, event: TEvent
+    ) -> Union[Tuple[Tensor, Tensor, Tensor], Tuple[Tensor, Tensor]]:
+        context_featurized, actions_featurized, selected = self.featurizer.featurize(
+            event
+        )
+
+        if len(context_featurized.dense) > 0:
+            raise NotImplementedError(
+                "pytorch policy doesn't support context with dense features"
+            )
+
+        for action_featurized in actions_featurized:
+            if len(action_featurized.dense) > 0:
+                raise NotImplementedError(
+                    "pytorch policy doesn't support action with dense features"
+                )
+
+        context_sparse = self.encode(
+            [self.convert_features_to_text(context_featurized.sparse)]
+        )
+
+        actions_sparse = []
+        for action_featurized in actions_featurized:
+            actions_sparse.append(
+                self.convert_features_to_text(action_featurized.sparse)
+            )
+        actions_sparse = self.encode(actions_sparse).unsqueeze(0)
+
+        if selected.score is not None:
+            return (
+                torch.Tensor([[selected.score]]),
+                context_sparse,
+                actions_sparse[:, selected.index, :].unsqueeze(1),
+            )
+        else:
+            return context_sparse, actions_sparse

src/learn_to_pick/pytorch/igw.py

@@ -0,0 +1,21 @@
+import torch
+from torch import Tensor
+from typing import Tuple
+
+
+def IGW(fhat: torch.Tensor, gamma: float) -> Tuple[Tensor, Tensor]:
+    from math import sqrt
+
+    fhatahat, ahat = fhat.max(dim=1)
+    A = fhat.shape[1]
+    gamma *= sqrt(A)
+    p = 1 / (A + gamma * (fhatahat.unsqueeze(1) - fhat))
+    sump = p.sum(dim=1)
+    p[range(p.shape[0]), ahat] += torch.clamp(1 - sump, min=0, max=None)
+    return torch.multinomial(p, num_samples=1).squeeze(1), ahat
+
+
+def SamplingIGW(A: Tensor, P: Tensor, gamma: float) -> list:
+    exploreind, _ = IGW(P, gamma)
+    explore = [ind for _, ind in zip(A, exploreind)]
+    return explore