File size: 3,931 Bytes
74e8f2f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 |
# Copyright 2024 Big Vision Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Utilities for generating zero-shot prompts."""
import re
import string
from typing import Sequence
from absl import logging
from big_vision.datasets.imagenet import class_names as imagenet_class_names
from big_vision.evaluators.proj.image_text import prompt_engineering_constants
import tensorflow_datasets as tfds
_CLASS_NAMES = { # For each dataset, maps from a source to its class names.
"imagenet2012": {
"clip": imagenet_class_names.CLIP_IMAGENET_CLASS_NAMES,
},
"grand-vision:imagenet2012": {
"clip": imagenet_class_names.CLIP_IMAGENET_CLASS_NAMES,
},
"imagenet_a": {
"clip": [
imagenet_class_names.CLIP_IMAGENET_CLASS_NAMES[i]
for i in imagenet_class_names.IMAGENET_A_LABELSET
]
},
"imagenet_r": {
"clip": [
imagenet_class_names.CLIP_IMAGENET_CLASS_NAMES[i]
for i in imagenet_class_names.IMAGENET_R_LABELSET
]
},
"imagenet_v2": {
"clip": imagenet_class_names.CLIP_IMAGENET_CLASS_NAMES,
},
}
_PROMPT_TEMPLATES = {
"class_name_only": ["{}"],
"clip_paper": prompt_engineering_constants.CLIP_PAPER_PROMPT_TEMPLATES,
"clip_best": prompt_engineering_constants.CLIP_BEST_PROMPT_TEMPLATES,
}
def get_class_names(*, dataset_name, source="dataset_info", canonicalize=True):
"""Returns class name for `dataset_name` from `source`."""
if isinstance(source, str):
if source.startswith("dataset_info:"):
name = source[len("dataset_info:"):]
class_names = tfds.builder(dataset_name).info.features[name].names
else:
class_names = _CLASS_NAMES[dataset_name][source]
else:
assert isinstance(source, Sequence) and all(
map(lambda s: isinstance(s, str), source)), source
class_names = source
if canonicalize:
class_names = [
canonicalize_text(name, keep_punctuation_exact_string=",")
for name in class_names
]
logging.info("Using %d class_names: %s", len(class_names), class_names)
return class_names
def get_prompt_templates(prompt_templates_name,
*,
canonicalize=True):
"""Returns prompt templates."""
prompts_templates = _PROMPT_TEMPLATES[prompt_templates_name]
if canonicalize:
prompts_templates = [
canonicalize_text(name, keep_punctuation_exact_string="{}")
for name in prompts_templates
]
logging.info("Using %d prompts_templates: %s", len(prompts_templates),
prompts_templates)
return prompts_templates
def canonicalize_text(text, *, keep_punctuation_exact_string=None):
"""Returns canonicalized `text` (lowercase and puncuation removed).
Args:
text: string to be canonicalized.
keep_punctuation_exact_string: If provided, then this exact string kept.
For example providing '{}' will keep any occurrences of '{}' (but will
still remove '{' and '}' that appear separately).
"""
text = text.replace("_", " ")
if keep_punctuation_exact_string:
text = keep_punctuation_exact_string.join(
part.translate(str.maketrans("", "", string.punctuation))
for part in text.split(keep_punctuation_exact_string))
else:
text = text.translate(str.maketrans("", "", string.punctuation))
text = text.lower()
text = re.sub(r"\s+", " ", text)
return text.strip()
|