Spaces:
Build error
Build error
File size: 3,722 Bytes
7025ca1 44a3896 7025ca1 44a3896 fd8b97e 022f82c fd8b97e 022f82c 44a3896 7025ca1 | 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 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | import numpy as np
from sentence_transformers import SentenceTransformer
from augmenator.keyword_catalog import AUGMENT_KEYWORDS, AugmentKeyword
MODEL_ID = "sentence-transformers/all-MiniLM-L6-v2"
SIMILARITY_THRESHOLD = 0.38
RELATIVE_MARGIN = 0.25
MAX_SELECTIONS = 5
MUTUAL_EXCLUSION_GROUPS: tuple[tuple[str, ...], ...] = (
("brighten", "darken"),
("contrast_up", "contrast_down"),
("warmer", "cooler"),
("saturate", "desaturate"),
("tint_red", "tint_green", "tint_blue"),
("cover_random", "cover_avoid_text", "cover_and_cutout", "cover_and_cutout_avoid_text"),
("add_cutout_transparent", "add_cutout_solid", "add_cutout", "add_cutout_avoid_text"),
(
"rotate",
"rotate_90_random",
"rotate_left",
"rotate_right",
"rotate_180",
"flip",
"flip_vertical",
),
("gamma_up", "gamma_down"),
(
"style_candy",
"style_mosaic",
"style_rain_princess",
"style_udnie",
"style_pointilism",
"style_starry_night",
"style_sketch",
"style_random",
),
)
_model: SentenceTransformer | None = None
_phrase_embeddings: dict[str, np.ndarray] = {}
def _load() -> SentenceTransformer:
global _model
if _model is None:
_model = SentenceTransformer(MODEL_ID)
return _model
def _build_phrase_cache(model: SentenceTransformer) -> None:
global _phrase_embeddings
if _phrase_embeddings:
return
for keyword in AUGMENT_KEYWORDS:
embeddings = model.encode(
list(keyword.phrases),
convert_to_numpy=True,
normalize_embeddings=True,
show_progress_bar=False,
)
_phrase_embeddings[keyword.id] = embeddings
def warmup() -> None:
model = _load()
_build_phrase_cache(model)
def _cosine_scores(instruction_embedding: np.ndarray) -> list[tuple[AugmentKeyword, float]]:
scored: list[tuple[AugmentKeyword, float]] = []
for keyword in AUGMENT_KEYWORDS:
phrase_embs = _phrase_embeddings[keyword.id]
similarities = phrase_embs @ instruction_embedding
score = float(np.max(similarities))
scored.append((keyword, score))
scored.sort(key=lambda item: item[1], reverse=True)
return scored
def _apply_mutual_exclusion(
matches: list[tuple[AugmentKeyword, float]],
) -> list[tuple[AugmentKeyword, float]]:
selected_ids: set[str] = set()
result: list[tuple[AugmentKeyword, float]] = []
for keyword, score in matches:
skip = False
for group in MUTUAL_EXCLUSION_GROUPS:
if keyword.id not in group:
continue
if any(member in selected_ids for member in group):
skip = True
break
if skip:
continue
selected_ids.add(keyword.id)
result.append((keyword, score))
return result
def score_keywords(instruction: str) -> list[tuple[AugmentKeyword, float]]:
warmup()
model = _load()
instruction_embedding = model.encode(
instruction,
convert_to_numpy=True,
normalize_embeddings=True,
show_progress_bar=False,
)
return _cosine_scores(instruction_embedding)
def select_keywords(instruction: str) -> list[tuple[AugmentKeyword, float]]:
scored = score_keywords(instruction)
if not scored:
return []
top_score = scored[0][1]
min_relative = top_score - RELATIVE_MARGIN
above_threshold = [
(keyword, score)
for keyword, score in scored
if score >= SIMILARITY_THRESHOLD and score >= min_relative
]
deduped = _apply_mutual_exclusion(above_threshold)
return deduped[:MAX_SELECTIONS]
|