crfm-helm 0.4.0__py3-none-any.whl → 0.5.1__py3-none-any.whl

helm/clients/bedrock_utils.py

@@ -0,0 +1,72 @@
+"""Helper utilities for working with Amazon Bedrock."""
+
+import os
+from typing import Optional
+
+from helm.common.hierarchical_logger import hlog
+from helm.common.optional_dependencies import handle_module_not_found_error
+
+try:
+    import boto3
+    from botocore.config import Config
+except ModuleNotFoundError as e:
+    handle_module_not_found_error(e, ["aws"])
+
+
+# From https://github.com/aws-samples/amazon-bedrock-workshop/blob/main/01_Generation/00_generate_w_bedrock.ipynb
+# MIT-0 Licensed
+def get_bedrock_client(
+    assumed_role: Optional[str] = None,
+    region: Optional[str] = None,
+    runtime: Optional[bool] = True,
+):
+    """Create a boto3 client for Amazon Bedrock, with optional configuration overrides
+
+    Parameters
+    ----------
+    assumed_role :
+        Optional ARN of an AWS IAM role to assume for calling the Bedrock service. If not
+        specified, the current active credentials will be used.
+    region :
+        Optional name of the AWS Region in which the service should be called (e.g. "us-east-1").
+        If not specified, AWS_REGION or AWS_DEFAULT_REGION environment variable will be used.
+    runtime :
+        Optional choice of getting different client to perform operations with the Amazon Bedrock service.
+    """
+    if region is None:
+        target_region = os.environ.get("AWS_REGION", os.environ.get("AWS_DEFAULT_REGION"))
+    else:
+        target_region = region
+
+    session_kwargs = {"region_name": target_region}
+    client_kwargs = {**session_kwargs}
+
+    profile_name = os.environ.get("AWS_PROFILE")
+    if profile_name:
+        session_kwargs["profile_name"] = profile_name
+
+    retry_config = Config(
+        region_name=target_region,
+        retries={
+            "max_attempts": 10,
+            "mode": "standard",
+        },
+    )
+    session = boto3.Session(**session_kwargs)
+
+    if assumed_role:
+        sts = session.client("sts")
+        response = sts.assume_role(RoleArn=str(assumed_role), RoleSessionName="crfm-helm")
+        client_kwargs["aws_access_key_id"] = response["Credentials"]["AccessKeyId"]
+        client_kwargs["aws_secret_access_key"] = response["Credentials"]["SecretAccessKey"]
+        client_kwargs["aws_session_token"] = response["Credentials"]["SessionToken"]
+
+    if runtime:
+        service_name = "bedrock-runtime"
+    else:
+        service_name = "bedrock"
+
+    bedrock_client = session.client(service_name=service_name, config=retry_config, **client_kwargs)
+
+    hlog(f"Amazon Bedrock client successfully created with endpoint {bedrock_client._endpoint}")
+    return bedrock_client

helm/{proxy/clients → clients}/client.py

@@ -1,11 +1,13 @@
 import json
 from abc import ABC, abstractmethod
-from typing import Dict, List, Optional
+from typing import List, Mapping, Optional, cast
 
 from helm.common.hierarchical_logger import hlog
 from helm.common.media_object import MultimediaObject, TEXT_TYPE
-from helm.common.request import Request, RequestResult, Sequence, Token
+from helm.common.request import Request, RequestResult, GeneratedOutput, Token
 from helm.common.cache import Cache, CacheConfig
+from helm.common.tokenization_request import DecodeRequest, TokenizationRequest
+from helm.tokenizers.tokenizer import Tokenizer
 
 
 class Client(ABC):
@@ -30,24 +32,28 @@ class CachingClient(Client):
         self.cache = Cache(cache_config) if cache_config is not None else None
 
     @staticmethod
-    def make_cache_key(raw_request: Dict, request: Request) -> Dict:
+    def make_cache_key(raw_request: Mapping, request: Request) -> Mapping:
         """
         Construct the key for the cache using the raw request.
         Add `request.random` to the key, if defined.
         """
         if request.random is not None:
             assert "random" not in raw_request
-            cache_key = {**raw_request, "random": request.random}
+            cache_key: Mapping = {**raw_request, "random": request.random}
         else:
             cache_key = raw_request
         return cache_key
 
 
-def truncate_sequence(sequence: Sequence, request: Request, print_warning: bool = True) -> Sequence:
+def truncate_sequence(sequence: GeneratedOutput, request: Request, print_warning: bool = True) -> GeneratedOutput:
     """
     Certain providers have bugs where they aren't respecting max_tokens,
     stop_sequences and the end of text token, so as a hack, we have to manually
     truncate the suffix of `sequence` and `tokens` as a post-hoc process.
+
+    This method is unsafe and may produce warnings or incorrect results.
+    Prefer using the safer truncate_and_tokenize_response_text() method instead
+    if your use case satisfies its requirements.
     """
     # TODO: if echo_prompt, then we should only ignore the prompt, but we don't
     # know how many tokens the prompt takes up.
@@ -87,7 +93,7 @@ def truncate_sequence(sequence: Sequence, request: Request, print_warning: bool
         if print_warning:
             hlog(f"WARNING: truncate_sequence needs to strip {json.dumps(stop)}")
 
-        sequence = Sequence(text=new_text, logprob=new_logprob, tokens=new_tokens)
+        sequence = GeneratedOutput(text=new_text, logprob=new_logprob, tokens=new_tokens)
 
     # Truncate based on the max number of tokens.
     if len(sequence.tokens) > request.max_tokens:
@@ -104,11 +110,63 @@ def truncate_sequence(sequence: Sequence, request: Request, print_warning: bool
 
         new_logprob = sum(token.logprob for token in new_tokens)
 
-        sequence = Sequence(text=new_text, logprob=new_logprob, tokens=new_tokens)
+        sequence = GeneratedOutput(text=new_text, logprob=new_logprob, tokens=new_tokens)
 
     return sequence
 
 
+def truncate_and_tokenize_response_text(
+    text: str, request: Request, tokenizer: Tokenizer, tokenizer_name: str, original_finish_reason: str = "endoftext"
+) -> GeneratedOutput:
+    """Truncate a string-only response to respect stop_sequences and max_tokens.
+
+    This can only be used if all of the following conditions are true:
+
+    - You have access to the tokenizer.
+    - The request has echo_prompt = False.
+    - The tokenizer supports encoding and decoding.
+    - The tokenizer's tokenize() method supports truncation.
+    - The model's response is text-only.
+    - The model's response not already provide the tokenized text.
+    - The model's response does not provide logprobs.
+
+    This method is safer than truncate_sequence() and should be preferred if the above conditions are met.
+    Unlike truncate_sequence(), this method will not produce warnings or incorrect results.
+    This is because the the tokens are derived from the truncated text using the tokenizer,
+    so the text and the tokens in the resulting result are guranteed to match."""
+    # Finish reason strings are token from basic_metrics._compute_finish_reason_metrics()
+    finish_reason: str = original_finish_reason
+    if request.echo_prompt:
+        raise Exception("truncate_and_tokenize_response_text() does not support requests with echo_prompt = True")
+
+    for stop_sequence in request.stop_sequences:
+        try:
+            text = text[: text.index(stop_sequence)]
+            finish_reason = "stop"
+        except ValueError:
+            pass
+
+    token_strings = cast(
+        List[str], tokenizer.tokenize(TokenizationRequest(text=text, tokenizer=tokenizer_name)).raw_tokens
+    )
+    if len(token_strings) > request.max_tokens:
+        encoded_ints = cast(
+            List[int],
+            tokenizer.tokenize(
+                TokenizationRequest(
+                    text=text, tokenizer=tokenizer_name, encode=True, truncation=True, max_length=request.max_tokens
+                )
+            ).raw_tokens,
+        )
+        text = tokenizer.decode(DecodeRequest(encoded_ints, tokenizer_name)).text
+        token_strings = cast(
+            List[str], tokenizer.tokenize(TokenizationRequest(text=text, tokenizer=tokenizer_name)).raw_tokens
+        )
+        finish_reason = "length"
+    tokens = [Token(text=token_string, logprob=0.0) for token_string in token_strings]
+    return GeneratedOutput(text=text, logprob=0.0, tokens=tokens, finish_reason={"reason": finish_reason})
+
+
 def cleanup_str(token: str, tokenizer_name: Optional[str] = None) -> str:
     """
     Certain tokenizers introduce special characters to represent spaces, such as

helm/clients/clip_score_client.py

@@ -0,0 +1,49 @@
+from typing import Dict, Optional
+from dataclasses import asdict
+
+from helm.common.cache import Cache, CacheConfig
+from helm.common.clip_score_request import DEFAULT_CLIP_SCORE_MODEL, CLIPScoreRequest, CLIPScoreResult
+from helm.clients.clip_scorers.base_clip_scorer import BaseCLIPScorer
+
+
+class CLIPScoreClientError(Exception):
+    pass
+
+
+class CLIPScoreClient:
+    def __init__(self, cache_config: CacheConfig):
+        self.cache = Cache(cache_config)
+        self._clip_scorer: Optional[BaseCLIPScorer] = None
+
+    def compute_score(self, request: CLIPScoreRequest) -> CLIPScoreResult:
+        """
+        Compute a CLIPScore for a given caption and image.
+        """
+        # TODO: support multilingual CLIPScore and other CLIP models.
+        assert request.model == DEFAULT_CLIP_SCORE_MODEL, f"Unsupported model: {request.model}"
+        assert not request.multilingual
+
+        try:
+
+            def do_it():
+                if self._clip_scorer is None:
+                    from helm.clients.clip_scorers.clip_scorer import CLIPScorer
+
+                    self._clip_scorer = CLIPScorer()
+
+                score: float = self._clip_scorer.compute_score(
+                    caption=request.caption, image_location=request.image_location
+                )
+                return {"score": score}
+
+            cache_key: Dict = asdict(request)
+            results, cached = self.cache.get(cache_key, do_it)
+
+        except Exception as e:
+            raise CLIPScoreClientError(e)
+
+        return CLIPScoreResult(
+            success=True,
+            cached=cached,
+            score=results["score"],
+        )

helm/clients/clip_scorers/base_clip_scorer.py

@@ -0,0 +1,18 @@
+from abc import abstractmethod, ABC
+from typing import List
+
+
+class BaseCLIPScorer(ABC):
+    @abstractmethod
+    def compute_score(self, caption: str, image_location: str) -> float:
+        pass
+
+    def select_best_image(self, caption: str, image_locations: List[str]) -> str:
+        """Selects the image from a list of images with the highest CLIPScore given the caption."""
+        assert len(image_locations) > 0, "Need at least one image"
+
+        if len(image_locations) == 1:
+            return image_locations[0]
+
+        scores: List[float] = [self.compute_score(caption, image_location) for image_location in image_locations]
+        return image_locations[scores.index(max(scores))]

helm/clients/clip_scorers/clip_scorer.py

@@ -0,0 +1,50 @@
+from typing import Literal
+
+from torchvision import transforms
+import torch
+
+from helm.common.gpu_utils import get_torch_device
+from helm.common.images_utils import open_image
+from helm.common.optional_dependencies import handle_module_not_found_error
+from .base_clip_scorer import BaseCLIPScorer
+
+
+_ = torch.manual_seed(42)
+
+
+class CLIPScorer(BaseCLIPScorer):
+    """
+    CLIPScore is a reference free metric that can be used to evaluate the correlation between an image
+    caption and the content of the image. It has been found to be highly correlated with human judgement.
+    Paper: https://arxiv.org/abs/2104.08718
+
+    We use the TorchMetrics implementation:
+    https://torchmetrics.readthedocs.io/en/stable/multimodal/clip_score.html.
+    The score is bound between 0 and 100, where a score closer to 100 is better.
+
+    Verified implementation against the scores of image-caption pairs from
+    https://wandb.ai/dalle-mini/dalle-mini/reports/OpenAI-CLIP-Score-exploration--VmlldzoxNjMwODM1.
+    """
+
+    def __init__(
+        self,
+        model_name: Literal[
+            "openai/clip-vit-base-patch16",
+            "openai/clip-vit-base-patch32",
+            "openai/clip-vit-large-patch14-336",
+            "openai/clip-vit-large-patch14",
+        ] = "openai/clip-vit-large-patch14",
+    ):
+        try:
+            from torchmetrics.multimodal import CLIPScore
+        except ModuleNotFoundError as e:
+            handle_module_not_found_error(e, ["heim"])
+
+        self._device: torch.device = get_torch_device()
+        self._metric = CLIPScore(model_name_or_path=model_name).to(self._device)
+
+    def compute_score(self, caption: str, image_location: str) -> float:
+        image = open_image(image_location)
+        image_tensor: torch.Tensor = transforms.ToTensor()(image).to(self._device)
+        score: float = self._metric(image_tensor, caption).detach().item()
+        return score

helm/clients/clip_scorers/multilingual_clip_scorer.py

@@ -0,0 +1,50 @@
+import torch
+import transformers
+
+from helm.common.gpu_utils import get_torch_device, get_torch_device_name
+from helm.common.images_utils import open_image
+from helm.common.optional_dependencies import handle_module_not_found_error
+from .base_clip_scorer import BaseCLIPScorer
+
+_ = torch.manual_seed(42)
+
+
+class MultilingualCLIPScorer(BaseCLIPScorer):
+    """
+    Multilingual-CLIP extends OpenAI's English text encoders to multiple other languages.
+    Adapted from https://huggingface.co/M-CLIP/XLM-Roberta-Large-Vit-L-14
+    """
+
+    TEXT_MODEL_NAME: str = "M-CLIP/XLM-Roberta-Large-Vit-L-14"
+    IMAGE_MODEL_NAME: str = "ViT-L/14"
+
+    def __init__(self):
+        try:
+            import clip
+            from multilingual_clip import pt_multilingual_clip
+        except ModuleNotFoundError as e:
+            handle_module_not_found_error(e, ["heim"])
+
+        super().__init__()
+        self._device: torch.device = get_torch_device()
+        self._text_model = pt_multilingual_clip.MultilingualCLIP.from_pretrained(self.TEXT_MODEL_NAME)
+        self._tokenizer = transformers.AutoTokenizer.from_pretrained(self.TEXT_MODEL_NAME)
+        self._model, self._preprocess = clip.load(self.IMAGE_MODEL_NAME, device=get_torch_device_name())
+
+    def compute_score(self, caption: str, image_location: str) -> float:
+        # Get text features
+        text_features = self._text_model.forward(caption, self._tokenizer)
+        text_features = text_features / text_features.norm(p=2, dim=-1, keepdim=True)
+        text_features = text_features.to(self._device)
+
+        image = open_image(image_location)
+        image = self._preprocess(image).unsqueeze(0).to(self._device)
+
+        # Get image features
+        with torch.no_grad():
+            image_features = self._model.encode_image(image)
+        image_features = image_features / image_features.norm(p=2, dim=-1, keepdim=True)
+
+        # Compute score using text and image features
+        score = 100 * (image_features * text_features).sum(axis=-1)
+        return score.detach().item()

helm/{proxy/clients → clients}/cohere_client.py

@@ -8,7 +8,7 @@ from helm.common.request import (
     EMBEDDING_UNAVAILABLE_REQUEST_RESULT,
     Request,
     RequestResult,
-    Sequence,
+    GeneratedOutput,
     Token,
 )
 from .client import CachingClient, truncate_sequence
@@ -120,7 +120,7 @@ class CohereClient(CachingClient):
             error: str = f"CohereClient error: {e}"
             return RequestResult(success=False, cached=False, error=error, completions=[], embedding=[])
 
-        completions: List[Sequence] = []
+        completions: List[GeneratedOutput] = []
         for generation in response["generations"]:
             # From https://docs.cohere.ai/generate-reference, "the likelihood refers to the average log-likelihood
             # of the entire specified string..." What we want is the sum of the log probabilities of all tokens.
@@ -132,14 +132,7 @@ class CohereClient(CachingClient):
                 logprob: float = token_likelihood.get("likelihood", 0)
                 sequence_logprob += logprob
 
-                tokens.append(
-                    Token(
-                        text=token_likelihood["token"],
-                        logprob=logprob,
-                        # Cohere does not include the top log probs in the response
-                        top_logprobs={},
-                    )
-                )
+                tokens.append(Token(text=token_likelihood["token"], logprob=logprob))
 
             sequence_text: str = generation["text"]
             if request.echo_prompt and request.max_tokens > 0:
@@ -147,7 +140,7 @@ class CohereClient(CachingClient):
                 # `return_likelihoods` is "ALL" and `max_tokens` is greater than 0.
                 sequence_text = request.prompt + sequence_text
 
-            completion: Sequence = Sequence(text=sequence_text, logprob=sequence_logprob, tokens=tokens)
+            completion: GeneratedOutput = GeneratedOutput(text=sequence_text, logprob=sequence_logprob, tokens=tokens)
             completion = truncate_sequence(completion, request)
             completions.append(completion)
 

helm/clients/gcs_client.py

@@ -0,0 +1,82 @@
+from dataclasses import asdict
+from typing import Dict, Optional
+import requests
+
+from helm.common.cache import Cache, CacheConfig
+from helm.common.optional_dependencies import handle_module_not_found_error
+from helm.common.hierarchical_logger import hlog
+from helm.common.file_upload_request import FileUploadRequest, FileUploadResult
+
+
+class GCSClientError(Exception):
+    pass
+
+
+class GCSClient:
+    """
+    Uploads files to GCS. Ensure the GOOGLE_APPLICATION_CREDENTIALS=/path/to/credentials.json
+    environment variable is set.
+    """
+
+    MAX_CHECK_ATTEMPTS: int = 10
+
+    def __init__(self, bucket_name: str, cache_config: CacheConfig):
+        try:
+            from google.cloud import storage  # type: ignore
+        except ModuleNotFoundError as e:
+            handle_module_not_found_error(e, ["heim"])
+
+        self._bucket_name: str = bucket_name
+        self._cache = Cache(cache_config)
+        self._storage_client: Optional[storage.Client] = None
+
+    def upload(self, request: FileUploadRequest) -> FileUploadResult:
+        """Uploads a file to GCS."""
+        try:
+            from google.cloud import storage  # type: ignore
+        except ModuleNotFoundError as e:
+            handle_module_not_found_error(e, ["heim"])
+
+        try:
+
+            def do_it():
+                if self._storage_client is None:
+                    self._storage_client = storage.Client()
+
+                bucket = self._storage_client.bucket(self._bucket_name)
+                file_path: str = request.path
+                blob = bucket.blob(file_path)
+
+                # Optional: set a generation-match precondition to avoid potential race conditions
+                # and data corruptions. The request to upload is aborted if the object's
+                # generation number does not match your precondition. For a destination
+                # object that does not yet exist, set the if_generation_match precondition to 0.
+                # If the destination object already exists in your bucket, set instead a
+                # generation-match precondition using its generation number.
+                generation_match_precondition: int = 0
+
+                blob.upload_from_filename(file_path, if_generation_match=generation_match_precondition)
+                url: str = self._get_url(file_path)
+
+                # Ensure the file was uploaded successfully
+                uploaded: bool = False
+                for _ in range(0, self.MAX_CHECK_ATTEMPTS):
+                    check_response = requests.head(url)
+                    if check_response.status_code == 200:
+                        uploaded = True
+                        break
+                assert uploaded, f"File {file_path} was not uploaded successfully."
+
+                hlog(f"File {file_path} uploaded and is available at {url}.")
+                return {"url": url}
+
+            cache_key: Dict = asdict(request)
+            result, cached = self._cache.get(cache_key, do_it)
+
+        except Exception as e:
+            raise GCSClientError(e)
+
+        return FileUploadResult(success=True, cached=cached, url=result["url"])
+
+    def _get_url(self, path: str) -> str:
+        return f"https://storage.googleapis.com/{self._bucket_name}/{path}"

helm/{proxy/clients → clients}/google_client.py

@@ -1,7 +1,7 @@
 from typing import List, Dict
 
 from helm.common.cache import CacheConfig
-from helm.common.request import Request, RequestResult, Sequence, Token
+from helm.common.request import Request, RequestResult, GeneratedOutput, Token
 from .client import CachingClient, truncate_sequence
 
 
@@ -32,7 +32,7 @@ class GoogleClient(CachingClient):
 
     def make_request(self, request: Request) -> RequestResult:
         raw_request = GoogleClient.convert_to_raw_request(request)
-        cache_key: Dict = CachingClient.make_cache_key(raw_request, request)
+        cache_key = CachingClient.make_cache_key(raw_request, request)
 
         try:
 
@@ -48,17 +48,17 @@ class GoogleClient(CachingClient):
             return RequestResult(success=False, cached=False, error=error, completions=[], embedding=[])
 
         # Expect the result to be structured the same way as a response from OpenAI API.
-        completions: List[Sequence] = []
+        completions: List[GeneratedOutput] = []
         for raw_completion in response["choices"]:
             sequence_logprob = 0
             tokens: List[Token] = []
 
             raw_data = raw_completion["logprobs"]
             for text, logprob in zip(raw_data["tokens"], raw_data["token_logprobs"]):
-                tokens.append(Token(text=text, logprob=logprob or 0, top_logprobs=dict()))
+                tokens.append(Token(text=text, logprob=logprob or 0))
                 sequence_logprob += logprob or 0
 
-            completion = Sequence(
+            completion = GeneratedOutput(
                 text=raw_completion["text"],
                 logprob=sequence_logprob,
                 tokens=tokens,

helm/clients/google_translate_client.py

@@ -0,0 +1,35 @@
+from typing import Optional
+
+from helm.common.cache import Cache, SqliteCacheConfig
+from helm.common.optional_dependencies import handle_module_not_found_error
+
+try:
+    from google.cloud import translate_v2 as translate  # type: ignore
+except ModuleNotFoundError as e:
+    handle_module_not_found_error(e, ["heim"])
+
+
+class GoogleTranslateClient:
+    """
+    Client for Google Translate.
+    Follow the instructions at https://cloud.google.com/translate/docs/setup to use this client.
+
+    # TODO: add this as a central service
+    """
+
+    def __init__(self, cache_path: str = "prod_env/cache/google_translate.sqlite"):
+        self.translate_client: Optional[translate.Client] = None
+        self.cache = Cache(SqliteCacheConfig(cache_path))
+
+    def translate(self, text: str, target_language: str) -> str:
+        def do_it():
+            if self.translate_client is None:
+                self.translate_client = translate.Client()
+
+            result = self.translate_client.translate(text, target_language=target_language)
+            del result["input"]
+            assert "translatedText" in result, f"Invalid response: {result}"
+            return result
+
+        response, _ = self.cache.get({"text": text, "target_language": target_language}, do_it)
+        return response["translatedText"]

helm/{proxy/clients → clients}/http_model_client.py

@@ -1,12 +1,13 @@
 import os
 from dataclasses import asdict
+from typing import Any, Dict
 
 from helm.common.cache import CacheConfig
 from helm.common.request import (
     wrap_request_time,
     Request,
     RequestResult,
-    Sequence,
+    GeneratedOutput,
     Token,
     EMBEDDING_UNAVAILABLE_REQUEST_RESULT,
 )
@@ -51,7 +52,7 @@ class HTTPModelClient(CachingClient):
 
         try:
 
-            def do_it():
+            def do_it() -> Dict[str, Any]:
                 url = f"{self.base_url}/process"
                 response = requests.post(url, json=raw_request, timeout=self.timeout)
                 response.raise_for_status()
@@ -63,11 +64,8 @@ class HTTPModelClient(CachingClient):
             else:
                 response, cached = do_it(), False
 
-            tokens = [
-                Token(text=token["text"], logprob=token["logprob"], top_logprobs=token["top_logprob"])
-                for token in response["tokens"]
-            ]
-            completions = [Sequence(text=response["text"], logprob=response["logprob"], tokens=tokens)]
+            tokens = [Token(text=token["text"], logprob=token["logprob"]) for token in response["tokens"]]
+            completions = [GeneratedOutput(text=response["text"], logprob=response["logprob"], tokens=tokens)]
 
             return RequestResult(
                 success=True,