crfm-helm 0.3.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

helm/clients/image_generation/mindalle/utils/sampling.py

@@ -0,0 +1,149 @@
+# ------------------------------------------------------------------------------------
+# minDALL-E
+# Copyright (c) 2021 Kakao Brain Corp. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------------------
+
+import torch
+from typing import Optional
+from tqdm import tqdm
+from torch.nn import functional as F
+
+
+def cutoff_topk_logits(logits: torch.FloatTensor, k: int) -> torch.FloatTensor:
+    if k is None:
+        return logits
+    else:
+        v, ix = torch.topk(logits, k)
+        out = logits.clone()
+        out[out < v[:, [-1]]] = -float("Inf")
+        return out
+
+
+def cutoff_topp_probs(probs: torch.FloatTensor, p: float) -> torch.FloatTensor:
+    if p is None:
+        return probs
+    else:
+        sorted_probs, sorted_indices = torch.sort(probs, dim=-1, descending=True)
+        cum_probs = torch.cumsum(sorted_probs, dim=-1)
+
+        sorted_idx_remove_cond = cum_probs >= p
+
+        sorted_idx_remove_cond[..., 1:] = sorted_idx_remove_cond[..., :-1].clone()
+        sorted_idx_remove_cond[..., 0] = 0
+
+        indices_to_remove = sorted_idx_remove_cond.scatter(-1, sorted_indices, sorted_idx_remove_cond)
+        probs = probs.masked_fill(indices_to_remove, 0.0)
+        norm_probs = probs / torch.sum(probs, dim=-1, keepdim=True)
+        return norm_probs
+
+
+def get_positional_encoding(inputs: torch.LongTensor, mode: str = "1d") -> torch.LongTensor:
+    device = inputs.device
+    if mode == "1d":
+        B, N = inputs.shape
+        xs_pos = torch.arange(N, device=device).repeat((B, 1))
+    elif mode == "2d":
+        B, H, W = inputs.shape
+        xs_pos_h = torch.arange(H, device=device).repeat(B, W, 1).transpose(1, 2)
+        xs_pos_w = torch.arange(W, device=device).repeat(B, H, 1)
+        xs_pos = (xs_pos_h, xs_pos_w)
+    else:
+        raise ValueError("%s positional encoding invalid" % mode)
+    return xs_pos
+
+
+@torch.no_grad()
+def sampling(
+    model: torch.nn.Module,
+    tokens: torch.LongTensor,
+    top_k: Optional[float] = None,
+    top_p: Optional[float] = None,
+    softmax_temperature: float = 1.0,
+    is_tqdm: bool = True,
+    use_fp16: bool = True,
+    max_seq_len: int = 256,
+) -> torch.LongTensor:
+    code = None
+    past = None
+
+    pbar = tqdm(range(max_seq_len), total=max_seq_len) if is_tqdm else range(max_seq_len)
+    pos_enc_tokens = get_positional_encoding(tokens, mode="1d")
+
+    for cnt, h in enumerate(pbar):
+        if code is None:
+            code_ = None
+            pos_enc_code_ = None
+        else:
+            code_ = code.clone().detach()
+            pos_enc_code_ = get_positional_encoding(code_, mode="1d")
+            code_ = code_[:, cnt - 1].unsqueeze(-1)
+            pos_enc_code_ = pos_enc_code_[:, cnt - 1].unsqueeze(-1)
+
+        logits, present = model.sampling(
+            images=code_, texts=tokens, pos_images=pos_enc_code_, pos_texts=pos_enc_tokens, use_fp16=use_fp16, past=past
+        )
+        logits = logits.to(dtype=torch.float32)
+        logits = logits / softmax_temperature
+
+        present = torch.stack(present).clone().detach()
+        if past is None:
+            past = [present]
+        else:
+            past.append(present)
+
+        logits = cutoff_topk_logits(logits, top_k)
+        probs = F.softmax(logits, dim=-1)
+        probs = cutoff_topp_probs(probs, top_p)
+
+        idx = torch.multinomial(probs, num_samples=1).clone().detach()
+        code = idx if code is None else torch.cat([code, idx], axis=1)
+
+    del past
+    return code
+
+
+@torch.no_grad()
+def sampling_igpt(
+    model: torch.nn.Module,
+    sos: torch.FloatTensor,
+    top_k: Optional[float] = None,
+    top_p: Optional[float] = None,
+    softmax_temperature: float = 1.0,
+    is_tqdm: bool = True,
+    use_fp16: bool = True,
+    max_seq_len: int = 256,
+) -> torch.LongTensor:
+    code = None
+    past = None
+    pbar = tqdm(range(max_seq_len), total=max_seq_len) if is_tqdm else range(max_seq_len)
+
+    for cnt, h in enumerate(pbar):
+        if code is None:
+            code_ = None
+            pos_enc_code_ = None
+        else:
+            code_ = code.clone().detach()
+            pos_enc_code_ = get_positional_encoding(code_, mode="1d")
+            code_ = code_[:, cnt - 1].unsqueeze(-1)
+            pos_enc_code_ = pos_enc_code_[:, cnt - 1].unsqueeze(-1)
+
+        logits, present = model.sampling(sos=sos, codes=code_, pos_codes=pos_enc_code_, use_fp16=use_fp16, past=past)
+        logits = logits.to(dtype=torch.float32)
+        logits = logits / softmax_temperature
+
+        present = torch.stack(present).clone().detach()
+        if past is None:
+            past = [present]
+        else:
+            past.append(present)
+
+        logits = cutoff_topk_logits(logits, top_k)
+        probs = F.softmax(logits, dim=-1)
+        probs = cutoff_topp_probs(probs, top_p)
+
+        idx = torch.multinomial(probs, num_samples=1).clone().detach()
+        code = idx if code is None else torch.cat([code, idx], axis=1)
+
+    del past
+    return code

helm/clients/image_generation/mindalle/utils/utils.py

@@ -0,0 +1,89 @@
+# ------------------------------------------------------------------------------------
+# minDALL-E
+# Copyright (c) 2021 Kakao Brain Corp. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------------------
+
+import os
+import random
+import urllib
+import hashlib
+import tarfile
+import torch
+import numpy as np
+from torch.nn import functional as F
+from tqdm import tqdm
+
+from helm.common.optional_dependencies import handle_module_not_found_error
+
+
+def set_seed(seed: int):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+@torch.no_grad()
+def clip_score(
+    prompt: str, images: np.ndarray, model_clip: torch.nn.Module, preprocess_clip, device: str
+) -> np.ndarray:
+    try:
+        import clip
+        from PIL import Image
+    except ModuleNotFoundError as e:
+        handle_module_not_found_error(e, ["heim"])
+
+    images = [preprocess_clip(Image.fromarray((image * 255).astype(np.uint8))) for image in images]
+    images = torch.stack(images, dim=0).to(device=device)
+    texts = clip.tokenize(prompt).to(device=device)
+    texts = torch.repeat_interleave(texts, images.shape[0], dim=0)
+
+    image_features = model_clip.encode_image(images)
+    text_features = model_clip.encode_text(texts)
+
+    scores = F.cosine_similarity(image_features, text_features).squeeze()
+    rank = torch.argsort(scores, descending=True).cpu().numpy()
+    return rank
+
+
+def download(url: str, root: str) -> str:
+    os.makedirs(root, exist_ok=True)
+    filename = os.path.basename(url)
+    pathname = filename[: -len(".tar.gz")]
+
+    expected_md5 = url.split("/")[-2]
+    download_target = os.path.join(root, filename)
+    result_path = os.path.join(root, pathname)
+
+    if os.path.isfile(download_target) and (os.path.exists(result_path) and not os.path.isfile(result_path)):
+        return result_path
+
+    with urllib.request.urlopen(url) as source, open(download_target, "wb") as output:
+        with tqdm(
+            total=int(source.info().get("Content-Length")), ncols=80, unit="iB", unit_scale=True, unit_divisor=1024
+        ) as loop:
+            while True:
+                buffer = source.read(8192)
+                if not buffer:
+                    break
+
+                output.write(buffer)
+                loop.update(len(buffer))
+
+    if hashlib.md5(open(download_target, "rb").read()).hexdigest() != expected_md5:
+        raise RuntimeError(f"Model has been downloaded but the md5 checksum does not not match")
+
+    with tarfile.open(download_target, "r:gz") as f:
+        pbar = tqdm(f.getmembers(), total=len(f.getmembers()))
+        for member in pbar:
+            pbar.set_description(f"extracting: {member.name} (size:{member.size // (1024 * 1024)}MB)")
+            f.extract(member=member, path=root)
+
+    return result_path
+
+
+def realpath_url_or_path(url_or_path: str, root: str = None) -> str:
+    if urllib.parse.urlparse(url_or_path).scheme in ("http", "https"):
+        return download(url_or_path, root)
+    return url_or_path

helm/clients/image_generation/mindalle_client.py

@@ -0,0 +1,115 @@
+from typing import Any, Dict, List
+
+import numpy as np
+
+from helm.common.cache import CacheConfig, Cache
+from helm.common.file_caches.file_cache import FileCache
+from helm.common.gpu_utils import get_torch_device_name
+from helm.common.hierarchical_logger import hlog, htrack_block
+from helm.common.optional_dependencies import handle_module_not_found_error
+from helm.common.request import Request, RequestResult, GeneratedOutput, wrap_request_time
+from helm.common.tokenization_request import (
+    DecodeRequest,
+    DecodeRequestResult,
+    TokenizationRequest,
+    TokenizationRequestResult,
+)
+from helm.clients.client import Client, CachingClient
+from .image_generation_client_utils import get_single_image_multimedia_object
+
+try:
+    from PIL import Image
+except ModuleNotFoundError as e:
+    handle_module_not_found_error(e, ["heim"])
+
+
+class MinDALLEClient(Client):
+    """
+    Source: https://github.com/kakaobrain/mindall-e
+    """
+
+    def __init__(self, cache_config: CacheConfig, file_cache: FileCache):
+        self._cache = Cache(cache_config)
+        self._file_cache: FileCache = file_cache
+
+        self._model = None
+
+    def _get_model(self):
+        try:
+            from helm.clients.image_generation.mindalle.models import Dalle
+        except ModuleNotFoundError as e:
+            handle_module_not_found_error(e, ["heim"])
+
+        if self._model is None:
+            self._model = Dalle.from_pretrained("minDALL-E/1.3B")
+            self._model = self._model.to(get_torch_device_name())
+        return self._model
+
+    def make_request(self, request: Request) -> RequestResult:
+        raw_request = {
+            "prompt": request.prompt,
+            # Setting this to a higher value can cause CUDA OOM
+            # Fix it to 1 and generate an image `request.num_completions` times
+            "num_candidates": 1,
+            "softmax_temperature": 1.0,
+            "top_k": 256,  # It is recommended that top_k is set lower than 256.
+            "top_p": None,
+            "device": "cuda",
+        }
+
+        try:
+
+            def do_it() -> Dict[str, Any]:
+                prompt: str = request.prompt
+
+                with htrack_block(f"Generating images for prompt: {prompt}"):
+                    model = self._get_model()
+
+                    images: List[Image] = []
+                    for _ in range(request.num_completions):
+                        output = model.sampling(**raw_request).cpu().numpy()
+                        output = np.transpose(output, (0, 2, 3, 1))
+                        image = Image.fromarray(np.asarray(output[0] * 255, dtype=np.uint8))
+                        images.append(image)
+
+                    assert (
+                        len(images) == request.num_completions
+                    ), f"Expected {request.num_completions} images, but got {len(images)}"
+
+                    result = {"file_locations": []}
+                    for image in images:
+                        # Write out the image to a file and save the path
+                        file_location: str = self._file_cache.get_unique_file_location()
+                        image.save(file_location)
+                        hlog(f"Image saved at {file_location}.")
+                        result["file_locations"].append(file_location)
+                    return result
+
+            # Include the model name and number of completions in the cache key
+            cache_key: Dict = CachingClient.make_cache_key(
+                {"model": request.model_engine, "n": request.num_completions, **raw_request}, request
+            )
+            results, cached = self._cache.get(cache_key, wrap_request_time(do_it))
+        except RuntimeError as ex:
+            error: str = f"MinDALLEClient error: {ex}"
+            return RequestResult(success=False, cached=False, error=error, completions=[], embedding=[])
+
+        completions: List[GeneratedOutput] = [
+            GeneratedOutput(
+                text="", logprob=0, tokens=[], multimodal_content=get_single_image_multimedia_object(location)
+            )
+            for location in results["file_locations"]
+        ]
+        return RequestResult(
+            success=True,
+            cached=cached,
+            request_time=results["request_time"],
+            completions=completions,
+            embedding=[],
+        )
+
+    def tokenize(self, request: TokenizationRequest) -> TokenizationRequestResult:
+        raise NotImplementedError("This client does not support tokenizing.")
+
+    def decode(self, request: DecodeRequest) -> DecodeRequestResult:
+        raise NotImplementedError("This client does not support decoding.")

helm/clients/image_generation/nudity_check_client.py

@@ -0,0 +1,64 @@
+from typing import Any, Dict, Optional
+import os
+
+from helm.common.cache import Cache, CacheConfig
+from helm.common.general import ensure_file_downloaded, ensure_directory_exists
+from helm.common.optional_dependencies import handle_module_not_found_error
+from helm.common.nudity_check_request import NudityCheckRequest, NudityCheckResult
+
+
+class NudityCheckClientError(Exception):
+    pass
+
+
+class NudityCheckClient:
+    MODEL_DOWNLOAD_URL: str = (
+        "https://worksheets.codalab.org/rest/bundles/0x81f7d1febb374fdcad6d9ae970f46750/contents/blob/"
+    )
+
+    def __init__(self, cache_config: CacheConfig):
+        try:
+            from nudenet import NudeClassifier
+        except ModuleNotFoundError as e:
+            handle_module_not_found_error(e, ["heim"])
+
+        self.cache = Cache(cache_config)
+        self._nudity_classifier: Optional[NudeClassifier] = None
+
+    def check_nudity(self, request: NudityCheckRequest) -> NudityCheckResult:
+        """Check for nudity for a batch of images using NudeNet."""
+        try:
+            from nudenet import NudeClassifier
+        except ModuleNotFoundError as e:
+            handle_module_not_found_error(e, ["heim"])
+
+        try:
+
+            def do_it() -> Dict[str, Any]:
+                if self._nudity_classifier is None:
+                    # The NudeNet library does not automatically provide model weights that work, so
+                    # manually download them. The path is hardcoded in the NudeNet library.
+                    base_path: str = os.path.join(os.path.expanduser("~"), ".NudeNet")
+                    ensure_directory_exists(base_path)
+                    model_path: str = os.path.join(base_path, "classifier_model.onnx")
+                    ensure_file_downloaded(source_url=self.MODEL_DOWNLOAD_URL, target_path=model_path)
+                    self._nudity_classifier = NudeClassifier()
+
+                path_to_nudity_scores: Dict[str, Dict[str, float]] = self._nudity_classifier.classify(
+                    request.image_locations
+                )
+                return path_to_nudity_scores
+
+            results, cached = self.cache.get({"locations": sorted(request.image_locations)}, do_it)
+        except Exception as e:
+            raise NudityCheckClientError(e)
+
+        nudity_results: Dict[str, bool] = {
+            image_location: nudity_result["unsafe"] > nudity_result["safe"]
+            for image_location, nudity_result in results.items()
+        }
+        return NudityCheckResult(
+            success=True,
+            cached=cached,
+            image_to_nudity=nudity_results,
+        )

helm/clients/image_generation/together_image_generation_client.py

@@ -0,0 +1,111 @@
+from typing import Any, Dict, List, Optional
+import base64
+import requests
+
+from helm.common.cache import CacheConfig, Cache
+from helm.common.file_caches.file_cache import FileCache
+from helm.common.request import Request, RequestResult, GeneratedOutput, wrap_request_time
+from helm.common.tokenization_request import (
+    TokenizationRequest,
+    TokenizationRequestResult,
+    DecodeRequest,
+    DecodeRequestResult,
+)
+
+from helm.clients.client import CachingClient, Client
+from .image_generation_client_utils import get_single_image_multimedia_object
+
+
+class TogetherImageGenerationClient(Client):
+    """
+    Client for image generation via the Together API.
+    """
+
+    DEFAULT_IMAGE_HEIGHT: int = 512
+    DEFAULT_IMAGE_WIDTH: int = 512
+
+    DEFAULT_GUIDANCE_SCALE: float = 7.5
+    DEFAULT_STEPS: int = 50
+
+    INFERENCE_ENDPOINT: str = "https://api.together.xyz/api/inference"
+
+    def __init__(self, cache_config: CacheConfig, file_cache: FileCache, api_key: Optional[str] = None):
+        self._cache = Cache(cache_config)
+        self.file_cache: FileCache = file_cache
+
+        self._promptist_model = None
+        self._promptist_tokenizer = None
+
+        self.api_key: Optional[str] = api_key
+
+    def make_request(self, request: Request) -> RequestResult:
+        # Following https://docs.together.xyz/en/api
+        assert request.image_generation_parameters is not None
+        raw_request = {
+            "request_type": "image-model-inference",
+            "model": request.model_engine,
+            "prompt": request.prompt,
+            "n": request.num_completions,
+            "guidance_scale": (
+                request.image_generation_parameters.guidance_scale
+                if request.image_generation_parameters.guidance_scale is not None
+                else self.DEFAULT_GUIDANCE_SCALE
+            ),
+            "steps": (
+                request.image_generation_parameters.diffusion_denoising_steps
+                if request.image_generation_parameters.diffusion_denoising_steps is not None
+                else self.DEFAULT_STEPS
+            ),
+        }
+
+        if (
+            request.image_generation_parameters.output_image_width is None
+            or request.image_generation_parameters.output_image_height is None
+        ):
+            raw_request["width"] = self.DEFAULT_IMAGE_WIDTH
+            raw_request["height"] = self.DEFAULT_IMAGE_HEIGHT
+        else:
+            raw_request["width"] = request.image_generation_parameters.output_image_width
+            raw_request["height"] = request.image_generation_parameters.output_image_height
+
+        cache_key = CachingClient.make_cache_key(raw_request, request)
+
+        try:
+
+            def do_it() -> Dict[str, Any]:
+                result = requests.post(self.INFERENCE_ENDPOINT, json=raw_request).json()
+                assert "output" in result, f"Invalid response: {result} from prompt: {request.prompt}"
+
+                for choice in result["output"]["choices"]:
+                    # Write out the image to a file and save the path
+                    choice["file_path"] = self.file_cache.store(lambda: base64.b64decode(choice["image_base64"]))
+                    choice.pop("image_base64", None)
+                return result["output"]
+
+            response, cached = self._cache.get(cache_key, wrap_request_time(do_it))
+        except RuntimeError as e:
+            error: str = f"TogetherVisionClient error: {e}"
+            return RequestResult(success=False, cached=False, error=error, completions=[], embedding=[])
+
+        completions: List[GeneratedOutput] = [
+            GeneratedOutput(
+                text="",
+                logprob=0,
+                tokens=[],
+                multimodal_content=get_single_image_multimedia_object(choice["file_path"]),
+            )
+            for choice in response["choices"]
+        ]
+        return RequestResult(
+            success=True,
+            cached=cached,
+            request_time=response["request_time"],
+            completions=completions,
+            embedding=[],
+        )
+
+    def tokenize(self, request: TokenizationRequest) -> TokenizationRequestResult:
+        raise NotImplementedError("This client does not support tokenizing.")
+
+    def decode(self, request: DecodeRequest) -> DecodeRequestResult:
+        raise NotImplementedError("This client does not support decoding.")

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

@@ -9,8 +9,8 @@ import torch
 
 from helm.common.cache import CacheConfig
 from helm.common.optional_dependencies import OptionalDependencyNotInstalled
-from helm.common.request import Request, RequestResult, Sequence, Token
-from helm.proxy.tokenizers.tokenizer import Tokenizer
+from helm.common.request import Request, RequestResult, GeneratedOutput, Token
+from helm.tokenizers.tokenizer import Tokenizer
 
 from .client import CachingClient
 from .lit_gpt_generate import generate  # type: ignore
@@ -89,6 +89,7 @@ class LitGPTClient(CachingClient):
     def __init__(
         self,
         tokenizer: Tokenizer,
+        tokenizer_name: str,
         cache_config: CacheConfig,
         checkpoint_dir: Path = Path(""),
         precision: str = "bf16-true",
@@ -97,7 +98,8 @@ class LitGPTClient(CachingClient):
         strategy: str = "auto",
         quantize: Optional[QuantizationType] = None,
     ):
-        super().__init__(cache_config=cache_config, tokenizer=tokenizer)
+        super().__init__(cache_config=cache_config)
+        self.tokenizer = tokenizer
         lit_gpt = LitGPT(checkpoint_dir, precision, device, devices, strategy, quantize)
         self.model = lit_gpt.model
         self.fabric = lit_gpt.fabric
@@ -154,8 +156,8 @@ class LitGPTClient(CachingClient):
 
         generated_tokens = []
         for token in tokens:
-            generated_tokens.append(Token(text=tokenizer.decode(token), logprob=0, top_logprobs={}))
-        completions = [Sequence(text=output, logprob=0, tokens=generated_tokens)]
+            generated_tokens.append(Token(text=tokenizer.decode(token), logprob=0))
+        completions = [GeneratedOutput(text=output, logprob=0, tokens=generated_tokens)]
 
         return RequestResult(
             success=True,

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

@@ -2,21 +2,22 @@ import json
 import requests
 from typing import Any, Dict, List
 import traceback
+from helm.common.cache import CacheConfig
 
 from helm.common.request import (
     wrap_request_time,
     EMBEDDING_UNAVAILABLE_REQUEST_RESULT,
     Request,
     RequestResult,
-    Sequence,
+    GeneratedOutput,
     Token,
 )
 from helm.common.tokenization_request import TokenizationRequest
-from helm.proxy.clients.huggingface_client import HuggingFaceClient
-from helm.proxy.clients.client import CachingClient, truncate_sequence
+from helm.clients.client import CachingClient, truncate_sequence
+from helm.tokenizers.tokenizer import Tokenizer
 
 
-class MegatronClient(HuggingFaceClient):
+class MegatronClient(CachingClient):
     """Client for remote Megatron-LM server.
 
     This client expects an external Megatron-LM server to be run on localhost:5000. See the
@@ -25,6 +26,11 @@ class MegatronClient(HuggingFaceClient):
     https://github.com/NVIDIA/Megatron-LM#gpt-text-generation
     """
 
+    def __init__(self, tokenizer: Tokenizer, tokenizer_name: str, cache_config: CacheConfig):
+        super().__init__(cache_config=cache_config)
+        self.tokenizer = tokenizer
+        self.tokenizer_name = tokenizer_name
+
     def _send_request(self, raw_request: Dict[str, Any]) -> Dict[str, Any]:
         response = requests.request(
             method="PUT",
@@ -43,10 +49,10 @@ class MegatronClient(HuggingFaceClient):
         return out
 
     def _tokenize_response(self, text: str) -> List[Token]:
-        tokenized_text = self.tokenizer.tokenize(TokenizationRequest(text, tokenizer="huggingface/gpt2"))
+        tokenized_text = self.tokenizer.tokenize(TokenizationRequest(text, tokenizer=self.tokenizer_name))
 
         # TODO(tgale): Support logprobs.
-        tokens = [Token(text=str(token), logprob=0, top_logprobs={}) for token in tokenized_text.raw_tokens]
+        tokens = [Token(text=str(token), logprob=0) for token in tokenized_text.raw_tokens]
         return tokens
 
     def _make_request(self, request: Request) -> RequestResult:
@@ -81,7 +87,7 @@ class MegatronClient(HuggingFaceClient):
 
         # NOTE: Megatron returns the de-tokenized response. Re-tokenize.
         tokens = self._tokenize_response(generated_text)
-        completion = Sequence(text=generated_text, logprob=0, tokens=tokens)
+        completion = GeneratedOutput(text=generated_text, logprob=0, tokens=tokens)
         completion = truncate_sequence(completion, request, print_warning=True)
 
         return RequestResult(