Added Claude 3 models, new invert_dict util

Author: adivekar-utexas

src/synthesizrr/base/algorithm/bedrock.py

@@ -3,7 +3,8 @@
 import os, time, logging, sys, shutil, numpy as np, pandas as pd, gc, warnings, json
 from contextlib import contextmanager
 from synthesizrr.base.util import optional_dependency, set_param_from_alias, Parameters, get_default, safe_validate_arguments, \
-    accumulate, dispatch, str_format_args, format_exception_msg, any_item, retry, Log, remove_values, as_list
+    accumulate, dispatch, dispatch_executor, any_are_none, format_exception_msg, any_item, retry, Log, remove_values, as_list, \
+    stop_executor
 from synthesizrr.base.framework import Dataset
 from synthesizrr.base.framework.task.text_generation import GenerativeLM, Prompts, GENERATED_TEXTS_COL, TextGenerationParams, \
     TextGenerationParamsMapper
@@ -17,38 +18,150 @@
     import boto3
 
 
+    def call_claude_v1_v2(
+            bedrock,
+            model_name: str,
+            prompt: str,
+            max_tokens_to_sample: int,
+            temperature: Optional[float] = None,
+            top_k: Optional[int] = None,
+            top_p: Optional[float] = None,
+            stop_sequences: Optional[List[str]] = None,
+            **kwargs,
+    ) -> str:
+        assert any_are_none(top_k, top_p), f'At least one of top_k, top_p must be None'
+        bedrock_params = {
+            "prompt": prompt,
+            "max_tokens_to_sample": max_tokens_to_sample,
+        }
+        if top_p is not None and temperature is not None:
+            raise ValueError(f'Cannot specify both top_p and temperature; at most one must be specified.')
+
+        if top_k is not None:
+            assert isinstance(top_k, int)
+            bedrock_params["top_k"] = top_k
+        elif temperature is not None:
+            assert isinstance(temperature, (float, int)) and 0 <= temperature <= 1
+            bedrock_params["temperature"] = temperature
+        elif top_p is not None:
+            assert isinstance(top_p, (float, int)) and 0 <= top_p <= 1
+            bedrock_params["top_p"] = top_p
+
+        if stop_sequences is not None:
+            bedrock_params["stop_sequences"] = stop_sequences
+
+        response = bedrock.invoke_model(
+            body=json.dumps(bedrock_params),
+            modelId=model_name,
+            accept='application/json',
+            contentType='application/json',
+        )
+        response_body: Dict = json.loads(response.get('body').read())
+        return response_body.get('completion')
+
+
+    def call_claude_v3(
+            bedrock,
+            *,
+            model_name: str,
+            prompt: str,
+            max_tokens_to_sample: int,
+            temperature: Optional[float] = None,
+            system: Optional[str] = None,
+            top_k: Optional[int] = None,
+            top_p: Optional[float] = None,
+            stop_sequences: Optional[List[str]] = None,
+            **kwargs,
+    ) -> str:
+        assert any_are_none(top_k, top_p), f'At least one of top_k, top_p must be None'
+        bedrock_params = {
+            "anthropic_version": "bedrock-2023-05-31",
+            "max_tokens": max_tokens_to_sample,
+            "messages": [
+                {
+                    "role": "user",
+                    "content": prompt,
+                }
+            ],
+        }
+        if system is not None:
+            assert isinstance(system, str) and len(system) > 0
+            bedrock_params["system"] = system
+
+        if top_p is not None and temperature is not None:
+            raise ValueError(f'Cannot specify both top_p and temperature; at most one must be specified.')
+
+        if top_k is not None:
+            assert isinstance(top_k, int) and len(system) >= 1
+            bedrock_params["top_k"] = top_k
+        elif top_p is not None:
+            assert isinstance(top_p, (float, int)) and 0 <= top_p <= 1
+            bedrock_params["top_p"] = top_p
+        elif temperature is not None:
+            assert isinstance(temperature, (float, int)) and 0 <= temperature <= 1
+            bedrock_params["temperature"] = temperature
+
+        if stop_sequences is not None:
+            bedrock_params["stop_sequences"] = stop_sequences
+
+        bedrock_params_json: str = json.dumps(bedrock_params)
+        # print(f'\n\nbedrock_params_json:\n{json.dumps(bedrock_params, indent=4)}')
+        response = bedrock.invoke_model(
+            body=bedrock_params_json,
+            modelId=model_name,
+            accept='application/json',
+            contentType='application/json',
+        )
+        response_body: Dict = json.loads(response.get('body').read())
+        return '\n'.join([d['text'] for d in response_body.get("content")])
+
+
     def call_bedrock(
             prompt: str,
             *,
             model_name: str,
             generation_params: Dict,
             region_name: List[str],
-    ) -> Dict:
-        start = time.perf_counter()
+    ) -> str:
         ## Note: creation of the bedrock client is fast.
         bedrock = boto3.client(
             service_name='bedrock-runtime',
             region_name=any_item(region_name),
-            # endpoint_url=f'https://bedrock.{region_name}.amazonaws.com',
+            # endpoint_url='https://bedrock.us-east-1.amazonaws.com',
         )
-        bedrock_invoke_model_params = {
-            "prompt": prompt,
-            **generation_params
-        }
-        response = bedrock.invoke_model(
-            body=json.dumps(bedrock_invoke_model_params),
-            modelId=model_name,
-            accept='application/json',
-            contentType='application/json'
-        )
-        response_body = json.loads(response.get('body').read())
-        end = time.perf_counter()
-        time_taken_sec: float = end - start
-        return response_body.get('completion')
+        if 'anthropic.claude-3' in model_name:
+            generated_text: str = call_claude_v3(
+                bedrock=bedrock,
+                prompt=prompt,
+                model_name=model_name,
+                **generation_params
+            )
+        elif 'claude' in model_name:
+            generated_text: str = call_claude_v1_v2(
+                bedrock=bedrock,
+                prompt=prompt,
+                model_name=model_name,
+                **generation_params
+            )
+        else:
+            bedrock_invoke_model_params = {
+                "prompt": prompt,
+                **generation_params
+            }
+            response = bedrock.invoke_model(
+                body=json.dumps(bedrock_invoke_model_params),
+                modelId=model_name,
+                accept='application/json',
+                contentType='application/json'
+            )
+            response_body = json.loads(response.get('body').read())
+            generated_text: str = response_body.get('completion')
+        return generated_text
 
 
     class BedrockPrompter(GenerativeLM):
         aliases = ['bedrock']
+        executor: Optional[Any] = None
 
         class Hyperparameters(GenerativeLM.Hyperparameters):
             ALLOWED_TEXT_GENERATION_PARAMS: ClassVar[List[str]] = [
@@ -59,6 +172,7 @@ class Hyperparameters(GenerativeLM.Hyperparameters):
                 'top_p',
                 'max_new_tokens',
                 'stop_sequences',
+                'system',
             ]
 
             region_name: List[str] = [
@@ -70,8 +184,9 @@ class Hyperparameters(GenerativeLM.Hyperparameters):
             model_name: constr(min_length=1)
             retries: conint(ge=0) = 3
             retry_wait: confloat(ge=0) = 1.0
-            retry_jitter: confloat(ge=0) = 0.25
+            retry_jitter: confloat(ge=0) = 0.5
             parallelize: Parallelize = Parallelize.sync
+            max_workers: int = 1
             generation_params: Union[TextGenerationParams, Dict, str]
 
             @root_validator(pre=True)
@@ -105,7 +220,15 @@ def max_num_generated_tokens(self) -> int:
 
         def initialize(self, model_dir: Optional[FileMetadata] = None):
             ## Ignore the model_dir.
-            pass
+            if self.executor is None:
+                self.executor: Optional[Any] = dispatch_executor(
+                    parallelize=self.hyperparams.parallelize,
+                    max_workers=self.hyperparams.max_workers,
+                )
+
+        def cleanup(self):
+            super(self.__class__, self).cleanup()
+            stop_executor(self.executor)
 
         @property
         def bedrock_text_generation_params(self) -> Dict[str, Any]:
@@ -146,6 +269,7 @@ def predict_step(self, batch: Prompts, **kwargs) -> Any:
                     self.prompt_model_with_retries,
                     prompt,
                     parallelize=self.hyperparams.parallelize,
+                    executor=self.executor,
                 )
                 generated_texts.append(generated_text)
             generated_texts: List[str] = accumulate(generated_texts)

src/synthesizrr/base/framework/mixins.py

@@ -939,19 +939,16 @@ def evaluate(
             self,
             metric: Optional[Union[Metric, Dict, str]] = None,
             *,
-            rolling: bool = False,
+            rolling: Optional[bool] = None,
+            inplace: Optional[bool] = None,
             **kwargs
     ) -> Metric:
         if metric is None:
-            return Metric.of(**kwargs).evaluate(self)
+            metric: Metric = Metric.of(**kwargs)
         if isinstance(metric, str):
-            return Metric.of(name=metric, **kwargs).evaluate(self)
-        if isinstance(metric, Metric):
-            if rolling:
-                return metric.evaluate(self, rolling=True)
-            else:
-                return metric.evaluate(self, inplace=False)
-        raise NotImplementedError(f'Unsupported value for input `metric`: {type(metric)} with value:\n{metric}')
+            metric: Metric = Metric.of(name=metric, **kwargs)
+        assert isinstance(metric, Metric)
+        return metric.evaluate(self, rolling=rolling, inplace=inplace)
 
     @safe_validate_arguments
     def columns(

src/synthesizrr/base/framework/task/text_generation.py

@@ -581,14 +581,14 @@ class BeamSearchParams(TextGenerationParams):
 
 class TopKSamplingParams(TextGenerationParams):
     strategy = 'TopKSampling'
-    temperature: confloat(gt=0.0, le=1.0)
+    temperature: confloat(gt=0.0, le=100.0) = 1.0
     do_sample: Literal[True] = True  ## When not doing greedy decoding, we should sample.
 
 
 class NucleusSamplingParams(TextGenerationParams):
     strategy = 'NucleusSampling'
     do_sample: Literal[True] = True  ## When not doing greedy decoding, we should sample.
-    temperature: confloat(gt=0.0, le=1.0)
+    temperature: confloat(gt=0.0, le=100.0) = 1.0
 
 
 class LogitsProcessorListParams(TextGenerationParams):

src/synthesizrr/base/util/language.py

@@ -1363,8 +1363,10 @@ def filter_keys(
     keys: Set = as_set(keys)
     if how == 'include':
         return keep_keys(d, keys)
-    else:
+    elif how == 'exclude':
         return remove_keys(d, keys)
+    else:
+        raise NotImplementedError(f'Invalid value for parameter `how`: "{how}"')
 
 
 def filter_values(
@@ -1488,6 +1490,15 @@ def eval_dict_values(params: Dict):
     return updated_dict
 
 
+def invert_dict(d: Dict) -> Dict:
+    if not isinstance(d, dict):
+        raise ValueError(f'{d} should be of type dict')
+    d_inv: Dict = {v: k for k, v in d.items()}
+    if len(d_inv) != len(d):
+        raise ValueError(f'Dict is not invertible as values are not unique.')
+    return d_inv
+
+
 ## ======================== NumPy utils ======================== ##
 def is_numpy_integer_array(data: Any) -> bool:
     if not isinstance(data, np.ndarray):
@@ -1929,6 +1940,10 @@ def iter_batches(
                 yield struct[i: min(i + batch_size, struct_len)]
 
 
+def mean(vals):
+    return sum(vals) / len(vals)
+
+
 def random_sample(
         data: Union[List, Tuple, np.ndarray],
         n: SampleSizeType,
@@ -2827,18 +2842,34 @@ class Timeout1Week(Timeout):
 
 
 @contextmanager
-def pd_display(
-        max_rows: Optional[int] = None,
-        max_cols: Optional[int] = None,
-        max_colwidth: Optional[int] = None,
-        vertical_align: str = 'top',
-        text_align: str = 'left',
-        ignore_css: bool = False,
-):
+def pd_display(**kwargs):
+    """
+    Use pd.describe_option('display') to see all options.
+    """
     try:
         from IPython.display import display
     except ImportError:
         display = print
+    set_param_from_alias(params=kwargs, param='max_rows', alias=['num_rows', 'nrows', 'rows'], default=None)
+    set_param_from_alias(params=kwargs, param='max_cols', alias=['num_cols', 'ncols', 'cols'], default=None)
+    set_param_from_alias(params=kwargs, param='max_colwidth', alias=[
+        'max_col_width',
+        'max_columnwidth', 'max_column_width',
+        'columnwidth', 'column_width',
+        'colwidth', 'col_width',
+    ], default=None)
+    set_param_from_alias(params=kwargs, param='vertical_align', alias=['valign'], default='top')
+    set_param_from_alias(params=kwargs, param='text_align', alias=['textalign'], default='left')
+    set_param_from_alias(params=kwargs, param='ignore_css', alias=['css'], default=False)
+
+    max_rows: Optional[int] = kwargs.get('max_rows')
+    max_cols: Optional[int] = kwargs.get('max_cols')
+    max_colwidth: Optional[int] = kwargs.get('max_colwidth')
+    vertical_align: str = kwargs['vertical_align']
+    text_align: str = kwargs['text_align']
+    ignore_css: bool = kwargs['ignore_css']
+
+    # print(kwargs)
 
     def disp(df: pd.DataFrame):
         css = [
@@ -2851,7 +2882,7 @@ def disp(df: pd.DataFrame):
                     ('padding', '10px'),
                 ]
             },
-            ## Align cell to top and left
+            ## Align cell to top and left/center
             {
                 'selector': 'td',
                 'props': [

src/synthesizrr/base/util/string.py

@@ -817,10 +817,10 @@ def is_fuzzy_match(cls, string: str, strings_to_match: List[str]) -> bool:
         return cls.fuzzy_match(string, strings_to_match) is not None
 
     @classmethod
-    def make_heading(cls, heading_text: str, width: int = 85, border: str = '=') -> str:
+    def header(cls, text: str, width: int = 65, border: str = '=') -> str:
         out = ''
         out += border * width + cls.NEWLINE
-        out += ('{:^' + str(width) + 's}').format(heading_text) + cls.NEWLINE
+        out += ('{:^' + str(width) + 's}').format(text) + cls.NEWLINE
         out += border * width + cls.NEWLINE
         return out