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

helm/clients/image_generation/dalle_mini/model/partitions.py

@@ -0,0 +1,84 @@
+import re
+
+from helm.common.optional_dependencies import handle_module_not_found_error
+
+try:
+    from flax.core.frozen_dict import freeze
+    from flax.traverse_util import flatten_dict, unflatten_dict
+    from jax.experimental import PartitionSpec as P
+except ModuleNotFoundError as e:
+    handle_module_not_found_error(e, ["heim"])
+
+
+# utils adapted from https://github.com/google-research/google-research/blob/master/flax_models/t5x/partitions.py
+# Sentinels
+_unmatched = object()
+
+# For specifying empty leaf dict `{}`
+empty_dict = object()
+
+
+def _match(qs, ks):
+    """Return True if regexes in qs match any window of strings in tuple ks."""
+    # compile regexes and force complete match
+    qts = tuple(map(lambda x: re.compile(x + "$"), qs))
+    for i in range(len(ks) - len(qs) + 1):
+        matches = [x.match(y) for x, y in zip(qts, ks[i:])]
+        if matches and all(matches):
+            return True
+    return False
+
+
+def _replacement_rules(rules):
+    def replace(key, val):
+        for rule, replacement in rules:
+            if _match(rule, key):
+                return replacement
+        return val
+
+    return replace
+
+
+def _get_partition_rules():
+    return [
+        # embeddings
+        (("embed_positions", "embedding"), P("mp", None)),
+        (("embed_tokens", "embedding"), P("mp", None)),
+        (("rel_bias", "embedding"), P(None, "mp")),
+        # attention
+        (("(q_proj|k_proj|v_proj)", "kernel"), P(None, "mp")),
+        (("out_proj", "kernel"), P("mp", None)),
+        # FFN
+        (("Dense_0", "kernel"), P(None, "mp")),
+        (("GLU.*", "Dense_1", "kernel"), P(None, "mp")),
+        (("GLU.*", "Dense_2", "kernel"), P("mp", None)),
+        (("FFN.*", "Dense_1", "kernel"), P("mp", None)),
+        # layer norms
+        (("(bias|scale)",), None),
+        (("lm_head", "kernel"), P(None, "mp")),
+        # head scale and tau
+        (("(head_scale|tau)",), None),
+    ]
+
+
+def set_partitions(in_dict, use_scan):
+    rules = _get_partition_rules()
+    replace = _replacement_rules(rules)
+    initd = {k: _unmatched for k in flatten_dict(in_dict)}
+    result = {k: replace(k, v) for k, v in initd.items()}
+    for k, v in result.items():
+        if v == _unmatched:
+            print(f"Unmatched -> {k}")
+    l = list(result.keys())
+    if use_scan:
+        # add None dimension to layers
+        result = {
+            k: (
+                (P(*(None,) + v) if v is not None else None)
+                if any(x in k for x in ["FlaxBartEncoderLayers", "FlaxBartDecoderLayers"])
+                else v
+            )
+            for k, v in result.items()
+        }
+    assert _unmatched not in result.values(), "Incomplete partition spec."
+    return freeze(unflatten_dict(result))

helm/clients/image_generation/dalle_mini/model/processor.py

@@ -0,0 +1,63 @@
+""" DalleBart processor """
+
+from typing import List
+
+from .configuration import DalleBartConfig
+from .text import TextNormalizer
+from .tokenizer import DalleBartTokenizer
+from .utils import PretrainedFromWandbMixin
+from helm.common.optional_dependencies import handle_module_not_found_error
+
+
+class DalleBartProcessorBase:
+    def __init__(self, tokenizer: DalleBartTokenizer, normalize_text: bool, max_text_length: int):
+        self.tokenizer = tokenizer
+        self.normalize_text = normalize_text
+        self.max_text_length = max_text_length
+        if normalize_text:
+            self.text_processor = TextNormalizer()
+        # create unconditional tokens
+        uncond = self.tokenizer(
+            "",
+            return_tensors="jax",
+            padding="max_length",
+            truncation=True,
+            max_length=self.max_text_length,
+        ).data
+        self.input_ids_uncond = uncond["input_ids"]
+        self.attention_mask_uncond = uncond["attention_mask"]
+
+    def __call__(self, text: List[str] = None):
+        try:
+            import jax.numpy as jnp
+        except ModuleNotFoundError as e:
+            handle_module_not_found_error(e, ["heim"])
+
+        # check that text is not a string
+        assert not isinstance(text, str), "text must be a list of strings"
+
+        if self.normalize_text:
+            text = [self.text_processor(t) for t in text]
+        res = self.tokenizer(
+            text,
+            return_tensors="jax",
+            padding="max_length",
+            truncation=True,
+            max_length=self.max_text_length,
+        ).data
+
+        # tokens used only with super conditioning
+        n = len(text)
+        res["input_ids_uncond"] = jnp.repeat(self.input_ids_uncond, n, axis=0)
+        res["attention_mask_uncond"] = jnp.repeat(self.attention_mask_uncond, n, axis=0)
+        return res
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        tokenizer = DalleBartTokenizer.from_pretrained(*args, **kwargs)
+        config = DalleBartConfig.from_pretrained(*args, **kwargs)
+        return cls(tokenizer, config.normalize_text, config.max_text_length)
+
+
+class DalleBartProcessor(PretrainedFromWandbMixin, DalleBartProcessorBase):
+    pass

helm/clients/image_generation/dalle_mini/model/text.py

@@ -0,0 +1,251 @@
+"""
+Utilities for processing text.
+"""
+
+import html
+import math
+import random
+import re
+from pathlib import Path
+
+import emoji
+from huggingface_hub import hf_hub_download
+
+from helm.common.optional_dependencies import handle_module_not_found_error
+
+try:
+    import ftfy
+    from unidecode import unidecode
+except ModuleNotFoundError as e:
+    handle_module_not_found_error(e, ["heim"])
+
+
+# based on wiki word occurrence
+person_token = [("a person", 282265), ("someone", 121194), ("somebody", 12219)]
+temp_token = "xtokx"  # avoid repeating chars
+
+
+class HashtagProcessor:
+    # Adapted from wordninja library
+    # We use our wikipedia word count + a good heuristic to make it work
+    def __init__(self):
+        wiki_word_frequency = hf_hub_download("dalle-mini/dalle-mini", filename="enwiki-words-frequency.txt")
+        self._word_cost = (l.split()[0] for l in Path(wiki_word_frequency).read_text(encoding="utf8").splitlines())
+        self._word_cost = {str(k): math.log(float(i + 1)) for i, k in enumerate(self._word_cost)}
+        self._max_word = max(len(x) for x in self._word_cost.keys())
+        self._SPLIT_RE = re.compile("[^a-zA-Z0-9']+")
+
+    def __call__(self, s):
+        """Uses dynamic programming to infer the location of spaces in a string without spaces."""
+        l = [self._split(x) for x in self._SPLIT_RE.split(s)]
+        return " ".join([item for sublist in l for item in sublist])
+
+    def _split(self, s):
+        # Find the best match for the i first characters, assuming cost has
+        # been built for the i-1 first characters.
+        # Returns a pair (match_cost, match_length).
+        def best_match(i):
+            candidates = enumerate(reversed(cost[max(0, i - self._max_word) : i]))
+            return min((c + self._word_cost.get(s[i - k - 1 : i].lower(), 9e999), k + 1) for k, c in candidates)
+
+        # Build the cost array
+        cost = [0]
+        for i in range(1, len(s) + 1):
+            c, k = best_match(i)
+            cost.append(c)
+
+        # Backtrack to recover the minimal-cost string.
+        out = []
+        i = len(s)
+        while i > 0:
+            c, k = best_match(i)
+            assert c == cost[i]
+            newToken = True
+            if not s[i - k : i] == "'":  # ignore a lone apostrophe
+                if len(out) > 0:
+                    # re-attach split 's and split digits
+                    if out[-1] == "'s" or (s[i - 1].isdigit() and out[-1][0].isdigit()):  # digit followed by digit
+                        out[-1] = s[i - k : i] + out[-1]  # combine current token with previous token
+                        newToken = False
+
+            if newToken:
+                out.append(s[i - k : i])
+
+            i -= k
+
+        return reversed(out)
+
+
+def replace_person_token(t):
+    "Used for CC12M"
+    t = re.sub("<person>([,\s]*(and)*[,\s]*<person>)+", " people ", t)
+    while "<person>" in t:
+        t = t.replace("<person>", f" {random.choices(*tuple(zip(*person_token)))[0]} ", 1)
+    return t
+
+
+def fix_html(t):
+    # from OpenAI CLIP
+    return html.unescape(html.unescape(t))
+
+
+def replace_punctuation_with_commas(t):
+    return re.sub("[()[\].,|:;?!=+~\-\/{}]", ",", t)
+
+
+def simplify_quotes(t):
+    return re.sub("""['"`]""", ' " ', t)
+
+
+def merge_quotes(t):
+    return re.sub('(\s*"+\s*)+', ' " ', t)
+
+
+def remove_comma_numbers(t):
+    def _f(t):
+        return re.sub("(\d),(\d{3})", r"\1\2", t)
+
+    return _f(_f(t))
+
+
+def pre_process_dot_numbers(t):
+    return re.sub("(\w)\.(\w)", rf"\1{temp_token}dot{temp_token}\2", t)
+
+
+def post_process_dot_numbers(t):
+    return re.sub(f"{temp_token}dot{temp_token}", ".", t)
+
+
+def pre_process_quotes(t):
+    # allows quotes only for 's, 't, 'd, 'm, 'll, 're, 've
+    return re.sub(r"'(?=([stdm]|(ll)|(re)|(ve)|(ll))\b)", rf"{temp_token}quote{temp_token}", t)
+
+
+def post_process_quotes(t):
+    return re.sub(f"{temp_token}quote{temp_token}", "'", t)
+
+
+def pre_process_dates(t):
+    return re.sub("(\d)/(\d)", rf"\1{temp_token}slash{temp_token}\2", t)
+
+
+def post_process_dates(t):
+    return re.sub(f"{temp_token}slash{temp_token}", "/", t)
+
+
+def merge_commas(t):
+    return re.sub("(\s*,+\s*)+", ", ", t)
+
+
+def add_space_after_commas(t):
+    return re.sub(",", ", ", t)
+
+
+def handle_special_chars(t):
+    "Handle special characters"
+    # replace "-" with a space when between words without space
+    t = re.sub("(\w)-(\w)", r"\1 \2", t)
+    # always add space around some characters
+    return re.sub("([%&\/$*])", r" \1 ", t)
+
+
+def expand_hashtags(t, hashtag_processor):
+    "Remove # and try to split words"
+    return re.sub("#(\w+)", lambda m: hashtag_processor(m.group(1)), t)
+
+
+_re_ignore_chars = r"[_#\\]"
+
+
+def ignore_chars(t):
+    "Ignore useless characters"
+    return re.sub(_re_ignore_chars, " ", t)
+
+
+def remove_extra_spaces(t):
+    "Remove extra spaces (including \t and \n)"
+    return re.sub("\s+", " ", t)
+
+
+def remove_repeating_chars(t):
+    "If the same character is present 4+ times (not 3 because of roman 'VIII'), replace with single instance"
+    return re.sub(r"(\D)(\1{3,})", r"\1", t)
+
+
+def remove_urls(t):
+    return re.sub(r"http\S+", "", t)
+
+
+def remove_html_tags(t):
+    return re.sub("<[^<]+?>", " ", t)
+
+
+def remove_first_last_commas(t):
+    t = t.strip()
+    t = t[:-1] if t and t[-1] == "," else t
+    t = t[1:] if t and t[0] == "," else t
+    return t.strip()
+
+
+def remove_wiki_ref(t):
+    t = re.sub(r"\A\s*\[\d+\]", "", t)
+    return re.sub(r"\[\d+\]\s*\Z", "", t)
+
+
+class TextNormalizer:
+    "Normalize text"
+
+    def __init__(self):
+        self._hashtag_processor = HashtagProcessor()
+
+    def __call__(self, t):
+        # fix some characters
+        t = ftfy.fix_text(t)
+        # fix html
+        t = fix_html(t)
+        # decode emojis (would be removed by unidecode)
+        t = emoji.demojize(t)
+        # decode and simplify text: see unidecode library
+        t = unidecode(t)
+        # lower case
+        t = t.lower()
+        # replace <PERSON> (for CC12M)
+        t = replace_person_token(t)
+        # remove wiki reference (for WIT)
+        t = remove_wiki_ref(t)
+        # remove html tags
+        t = remove_html_tags(t)
+        # remove urls
+        t = remove_urls(t)
+        # remove commas in numbers
+        t = remove_comma_numbers(t)
+        # handle dots in numbers and quotes - Part 1
+        t = pre_process_dot_numbers(t)
+        t = pre_process_quotes(t)
+        t = pre_process_dates(t)
+        # handle special characters
+        t = handle_special_chars(t)
+        # handle hashtags
+        t = expand_hashtags(t, self._hashtag_processor)
+        # ignore useless characters
+        t = ignore_chars(t)
+        # simplify quotes
+        t = simplify_quotes(t)
+        # all punctuation becomes commas
+        t = replace_punctuation_with_commas(t)
+        # handle dots in numbers and quotes - Part 2
+        t = post_process_dot_numbers(t)
+        t = post_process_quotes(t)
+        t = post_process_dates(t)
+        # handle repeating characters
+        t = remove_repeating_chars(t)
+        # merge quotes
+        t = merge_quotes(t)
+        # merge commas
+        t = merge_commas(t)
+        # remove multiple spaces
+        t = remove_extra_spaces(t)
+        # remove first and last comma
+        t = remove_first_last_commas(t)
+        # always start with a space
+        return f" {t}"

helm/clients/image_generation/dalle_mini/model/tokenizer.py

@@ -0,0 +1,9 @@
+""" DalleBart tokenizer """
+
+from transformers import BartTokenizerFast
+
+from .utils import PretrainedFromWandbMixin
+
+
+class DalleBartTokenizer(PretrainedFromWandbMixin, BartTokenizerFast):
+    pass

helm/clients/image_generation/dalle_mini/model/utils.py

@@ -0,0 +1,29 @@
+import os
+import tempfile
+
+from helm.common.optional_dependencies import handle_module_not_found_error
+
+
+class PretrainedFromWandbMixin:
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        """
+        Initializes from a wandb artifact or delegates loading to the superclass.
+        """
+        try:
+            import wandb
+        except ModuleNotFoundError as e:
+            handle_module_not_found_error(e, ["heim"])
+
+        with tempfile.TemporaryDirectory() as tmp_dir:  # avoid multiple artifact copies
+            if ":" in pretrained_model_name_or_path and not os.path.isdir(pretrained_model_name_or_path):
+                # wandb artifact
+                if wandb.run is not None:
+                    artifact = wandb.run.use_artifact(pretrained_model_name_or_path)
+                else:
+                    artifact = wandb.Api().artifact(pretrained_model_name_or_path)
+                pretrained_model_name_or_path = artifact.download(tmp_dir)
+
+            return super(PretrainedFromWandbMixin, cls).from_pretrained(
+                pretrained_model_name_or_path, *model_args, **kwargs
+            )

helm/clients/image_generation/dalle_mini/vqgan_jax/__init__.py

@@ -0,0 +1 @@
+from . import *

helm/clients/image_generation/dalle_mini/vqgan_jax/configuration_vqgan.py

@@ -0,0 +1,40 @@
+from typing import Tuple
+
+from transformers import PretrainedConfig
+
+
+class VQGANConfig(PretrainedConfig):
+    def __init__(
+        self,
+        ch: int = 128,
+        out_ch: int = 3,
+        in_channels: int = 3,
+        num_res_blocks: int = 2,
+        resolution: int = 256,
+        z_channels: int = 256,
+        ch_mult: Tuple = (1, 1, 2, 2, 4),
+        attn_resolutions: int = (16,),
+        n_embed: int = 1024,
+        embed_dim: int = 256,
+        dropout: float = 0.0,
+        double_z: bool = False,
+        resamp_with_conv: bool = True,
+        give_pre_end: bool = False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.ch = ch
+        self.out_ch = out_ch
+        self.in_channels = in_channels
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.z_channels = z_channels
+        self.ch_mult = list(ch_mult)
+        self.attn_resolutions = list(attn_resolutions)
+        self.n_embed = n_embed
+        self.embed_dim = embed_dim
+        self.dropout = dropout
+        self.double_z = double_z
+        self.resamp_with_conv = resamp_with_conv
+        self.give_pre_end = give_pre_end
+        self.num_resolutions = len(ch_mult)

helm/clients/image_generation/dalle_mini/vqgan_jax/convert_pt_model_to_jax.py

@@ -0,0 +1,107 @@
+import re
+
+import torch
+
+from .modeling_flax_vqgan import VQModel
+from .configuration_vqgan import VQGANConfig
+from helm.common.optional_dependencies import handle_module_not_found_error
+
+try:
+    import jax.numpy as jnp
+    from flax.traverse_util import flatten_dict, unflatten_dict
+except ModuleNotFoundError as e:
+    handle_module_not_found_error(e, ["heim"])
+
+
+regex = r"\w+[.]\d+"
+
+
+def rename_key(key):
+    pats = re.findall(regex, key)
+    for pat in pats:
+        key = key.replace(pat, "_".join(pat.split(".")))
+    return key
+
+
+# Adapted from https://github.com/huggingface/transformers/blob/ff5cdc086be1e0c3e2bbad8e3469b34cffb55a85/src/transformers/modeling_flax_pytorch_utils.py#L61
+def convert_pytorch_state_dict_to_flax(pt_state_dict, flax_model):
+    # convert pytorch tensor to numpy
+    pt_state_dict = {k: v.numpy() for k, v in pt_state_dict.items()}
+
+    random_flax_state_dict = flatten_dict(flax_model.params)
+    flax_state_dict = {}
+
+    remove_base_model_prefix = (flax_model.base_model_prefix not in flax_model.params) and (
+        flax_model.base_model_prefix in set([k.split(".")[0] for k in pt_state_dict.keys()])
+    )
+    add_base_model_prefix = (flax_model.base_model_prefix in flax_model.params) and (
+        flax_model.base_model_prefix not in set([k.split(".")[0] for k in pt_state_dict.keys()])
+    )
+
+    # Need to change some parameters name to match Flax names so that we don't have to fork any layer
+    for pt_key, pt_tensor in pt_state_dict.items():
+        pt_tuple_key = tuple(pt_key.split("."))
+
+        has_base_model_prefix = pt_tuple_key[0] == flax_model.base_model_prefix
+        require_base_model_prefix = (flax_model.base_model_prefix,) + pt_tuple_key in random_flax_state_dict
+
+        if remove_base_model_prefix and has_base_model_prefix:
+            pt_tuple_key = pt_tuple_key[1:]
+        elif add_base_model_prefix and require_base_model_prefix:
+            pt_tuple_key = (flax_model.base_model_prefix,) + pt_tuple_key
+
+        # Correctly rename weight parameters
+        if (
+            "norm" in pt_key
+            and (pt_tuple_key[-1] == "bias")
+            and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)
+            and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
+        ):
+            pt_tuple_key = pt_tuple_key[:-1] + ("scale",)
+        elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
+            pt_tuple_key = pt_tuple_key[:-1] + ("scale",)
+        if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
+            pt_tuple_key = pt_tuple_key[:-1] + ("embedding",)
+        elif pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4 and pt_tuple_key not in random_flax_state_dict:
+            # conv layer
+            pt_tuple_key = pt_tuple_key[:-1] + ("kernel",)
+            pt_tensor = pt_tensor.transpose(2, 3, 1, 0)
+        elif pt_tuple_key[-1] == "weight" and pt_tuple_key not in random_flax_state_dict:
+            # linear layer
+            pt_tuple_key = pt_tuple_key[:-1] + ("kernel",)
+            pt_tensor = pt_tensor.T
+        elif pt_tuple_key[-1] == "gamma":
+            pt_tuple_key = pt_tuple_key[:-1] + ("weight",)
+        elif pt_tuple_key[-1] == "beta":
+            pt_tuple_key = pt_tuple_key[:-1] + ("bias",)
+
+        if pt_tuple_key in random_flax_state_dict:
+            if pt_tensor.shape != random_flax_state_dict[pt_tuple_key].shape:
+                raise ValueError(
+                    f"PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape "
+                    f"{random_flax_state_dict[pt_tuple_key].shape}, but is {pt_tensor.shape}."
+                )
+
+        # also add unexpected weight so that warning is thrown
+        flax_state_dict[pt_tuple_key] = jnp.asarray(pt_tensor)
+
+    return unflatten_dict(flax_state_dict)
+
+
+def convert_model(config_path, pt_state_dict_path, save_path):
+    config = VQGANConfig.from_pretrained(config_path)
+    model = VQModel(config)
+
+    state_dict = torch.load(pt_state_dict_path, map_location="cpu")["state_dict"]
+    keys = list(state_dict.keys())
+    for key in keys:
+        if key.startswith("loss"):
+            state_dict.pop(key)
+            continue
+        renamed_key = rename_key(key)
+        state_dict[renamed_key] = state_dict.pop(key)
+
+    state = convert_pytorch_state_dict_to_flax(state_dict, model)
+    model.params = state
+    model.save_pretrained(save_path)
+    return model