Merge pull request #168 from generative-computing/hen/modular_sampling

fix: convert sampling to a package

Author: HendrikStrobelt

docs/examples/best_of_n/prm.py

@@ -8,7 +8,7 @@
 )
 from mellea.backends.types import ModelOption
 from mellea.stdlib.rewards.prm_scorer import PRMScorer
-from mellea.stdlib.sampling import BestofNSamplingStrategy
+from mellea.stdlib.sampling.best_of_n import BestofNSamplingStrategy
 
 # create a session for the generator using Granite 3.3 8B on Huggingface and a simple context [see below]
 m = start_session(backend_name="hf", model_options={ModelOption.MAX_NEW_TOKENS: 512})

mellea/stdlib/funcs.py

@@ -40,7 +40,7 @@ def act(
     context: Context,
     backend: Backend,
     *,
-    strategy: SamplingStrategy | None = RejectionSamplingStrategy(loop_budget=1),
+    strategy: SamplingStrategy | None = RejectionSamplingStrategy(loop_budget=2),
     return_sampling_results: Literal[False] = False,
     format: type[BaseModelSubclass] | None = None,
     model_options: dict | None = None,
@@ -54,7 +54,7 @@ def act(
     context: Context,
     backend: Backend,
     *,
-    strategy: SamplingStrategy | None = RejectionSamplingStrategy(loop_budget=1),
+    strategy: SamplingStrategy | None = RejectionSamplingStrategy(loop_budget=2),
     return_sampling_results: Literal[True],
     format: type[BaseModelSubclass] | None = None,
     model_options: dict | None = None,
@@ -68,7 +68,7 @@ def act(
     backend: Backend,
     *,
     requirements: list[Requirement] | None = None,
-    strategy: SamplingStrategy | None = RejectionSamplingStrategy(loop_budget=1),
+    strategy: SamplingStrategy | None = RejectionSamplingStrategy(loop_budget=2),
     return_sampling_results: bool = False,
     format: type[BaseModelSubclass] | None = None,
     model_options: dict | None = None,
@@ -114,7 +114,7 @@ async def _act(
     backend: Backend,
     *,
     requirements: list[Requirement] | None = None,
-    strategy: SamplingStrategy | None = RejectionSamplingStrategy(loop_budget=1),
+    strategy: SamplingStrategy | None = RejectionSamplingStrategy(loop_budget=2),
     return_sampling_results: bool = False,
     format: type[BaseModelSubclass] | None = None,
     model_options: dict | None = None,
@@ -213,7 +213,7 @@ def instruct(
     user_variables: dict[str, str] | None = None,
     prefix: str | CBlock | None = None,
     output_prefix: str | CBlock | None = None,
-    strategy: SamplingStrategy | None = RejectionSamplingStrategy(loop_budget=1),
+    strategy: SamplingStrategy | None = RejectionSamplingStrategy(loop_budget=2),
     return_sampling_results: Literal[False] = False,
     format: type[BaseModelSubclass] | None = None,
     model_options: dict | None = None,
@@ -234,7 +234,7 @@ def instruct(
     user_variables: dict[str, str] | None = None,
     prefix: str | CBlock | None = None,
     output_prefix: str | CBlock | None = None,
-    strategy: SamplingStrategy | None = RejectionSamplingStrategy(loop_budget=1),
+    strategy: SamplingStrategy | None = RejectionSamplingStrategy(loop_budget=2),
     return_sampling_results: Literal[True],
     format: type[BaseModelSubclass] | None = None,
     model_options: dict | None = None,
@@ -254,7 +254,7 @@ def instruct(
     user_variables: dict[str, str] | None = None,
     prefix: str | CBlock | None = None,
     output_prefix: str | CBlock | None = None,
-    strategy: SamplingStrategy | None = RejectionSamplingStrategy(loop_budget=1),
+    strategy: SamplingStrategy | None = RejectionSamplingStrategy(loop_budget=2),
     return_sampling_results: bool = False,
     format: type[BaseModelSubclass] | None = None,
     model_options: dict | None = None,