1199 lines
45 KiB
Python
1199 lines
45 KiB
Python
# ComfyUI-WordEmbeddings
|
|
# Core nodes with a typed WE_MODEL handle:
|
|
# - WordEmbeddingsLoader (pretrained dropdown -> WE_MODEL)
|
|
# - WordEmbeddingsLocalModelLoader (local file -> WE_MODEL)
|
|
# - WordEmbeddingsExplorer (WE_MODEL + word -> neighbors)
|
|
# - WordEmbeddingsInterpolator (WE_MODEL + word_a + word_b -> samples)
|
|
# - WordEmbeddingsEquation (WE_MODEL + expression -> analogy)
|
|
# - WordEmbeddingsTokenNeighbors (WE_MODEL + word -> neighbor words + cosines)
|
|
# - WordEmbeddingsTokenAxis (WE_MODEL + word (+pos,neg) -> projection & cosines + summary)
|
|
# - WordEmbeddingsTokenAxis2D (WE_MODEL + token + 2 axes -> x,y + summary)
|
|
# - WordEmbeddingsTokenAxis3D (WE_MODEL + token + 3 axes -> x,y,z + summary)
|
|
# - WordEmbeddingsTokenCentrality (WE_MODEL + word -> cosine to global mean & norm)
|
|
|
|
import os
|
|
import re
|
|
import math
|
|
import numpy as np
|
|
from dataclasses import dataclass
|
|
|
|
# ---------- gensim helpers ----------
|
|
def _lazy_api():
|
|
import gensim.downloader as api
|
|
return api
|
|
|
|
def _list_gensim_models():
|
|
SMALL_DEFAULT = "glove-wiki-gigaword-50"
|
|
try:
|
|
api = _lazy_api()
|
|
info = api.info()
|
|
names = list(info.get("models", {}).keys())
|
|
if not names:
|
|
return [SMALL_DEFAULT]
|
|
if SMALL_DEFAULT in names:
|
|
names.remove(SMALL_DEFAULT)
|
|
names.insert(0, SMALL_DEFAULT)
|
|
return names
|
|
except Exception:
|
|
return [SMALL_DEFAULT]
|
|
|
|
try:
|
|
from gensim.models import KeyedVectors
|
|
except Exception:
|
|
KeyedVectors = None
|
|
|
|
_MODEL_CHOICES = _list_gensim_models()
|
|
|
|
# ---------- WE_MODEL handle ----------
|
|
@dataclass
|
|
class WEModel:
|
|
name: str # model label (e.g., "glove-wiki-gigaword-50" or "local:foo.bin")
|
|
kv: object # gensim KeyedVectors or equivalent
|
|
|
|
def _load_pretrained(name: str):
|
|
api = _lazy_api()
|
|
return api.load(name)
|
|
|
|
def _load_local_keyedvectors(path: str, fmt: str = "auto"):
|
|
if KeyedVectors is None:
|
|
raise RuntimeError("gensim KeyedVectors not available (pip install gensim).")
|
|
p = path.strip().strip('"').strip("'")
|
|
if not os.path.exists(p):
|
|
raise FileNotFoundError(f"File not found: {p}")
|
|
|
|
if fmt == "kv":
|
|
return KeyedVectors.load(p, mmap="r")
|
|
if fmt == "word2vec_bin":
|
|
return KeyedVectors.load_word2vec_format(p, binary=True)
|
|
if fmt == "word2vec_text":
|
|
return KeyedVectors.load_word2vec_format(p, binary=False)
|
|
|
|
# auto detect
|
|
ext = os.path.splitext(p)[1].lower()
|
|
if ext in [".kv", ".kv2"]:
|
|
return KeyedVectors.load(p, mmap="r")
|
|
if ext in [".bin", ".bin.gz"]:
|
|
return KeyedVectors.load_word2vec_format(p, binary=True)
|
|
if ext in [".txt", ".vec", ".gz"]:
|
|
return KeyedVectors.load_word2vec_format(p, binary=False)
|
|
|
|
# last attempts
|
|
try:
|
|
return KeyedVectors.load(p, mmap="r")
|
|
except Exception:
|
|
return KeyedVectors.load_word2vec_format(p, binary=False)
|
|
|
|
# ---------- vector helpers ----------
|
|
def _kv_has_word(kv, w: str) -> bool:
|
|
if hasattr(kv, "key_to_index"):
|
|
return w in kv.key_to_index
|
|
try:
|
|
_ = kv[w]
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
def _kv_get_vec_case_insensitive(kv, w: str):
|
|
try:
|
|
return kv[w]
|
|
except KeyError:
|
|
lw = w.lower()
|
|
if _kv_has_word(kv, lw):
|
|
return kv[lw]
|
|
raise
|
|
|
|
def _cos(u, v):
|
|
nu, nv = np.linalg.norm(u), np.linalg.norm(v)
|
|
if nu == 0 or nv == 0:
|
|
return 0.0
|
|
return float((u @ v) / (nu * nv))
|
|
|
|
def _slerp(a, b, t):
|
|
a = a / np.linalg.norm(a); b = b / np.linalg.norm(b)
|
|
dot = float(np.clip(a @ b, -1.0, 1.0))
|
|
if dot > 0.9995:
|
|
v = (1 - t) * a + t * b
|
|
n = np.linalg.norm(v)
|
|
return v / n if n > 0 else v
|
|
theta = math.acos(dot)
|
|
return (math.sin((1 - t) * theta) * a + math.sin(t * theta) * b) / math.sin(theta)
|
|
|
|
def _nearest_to(v, kv, topn=1000, seen=None):
|
|
cand = kv.similar_by_vector(v, topn=topn)
|
|
best = None
|
|
best_sim = -1.0
|
|
for w, _ in cand:
|
|
if seen and w in seen:
|
|
continue
|
|
x = kv[w]
|
|
sim = _cos(v, x)
|
|
if sim > best_sim:
|
|
best_sim = sim
|
|
best = (w, x, best_sim)
|
|
return best
|
|
|
|
# ---------- caching for heavy global stats ----------
|
|
def _get_or_make_mean_vec(kv):
|
|
mv = getattr(kv, "_we_mean_vec", None)
|
|
if mv is not None:
|
|
return mv
|
|
try:
|
|
M = kv.get_normed_vectors() # (n_words, dim), rows L2-normalized
|
|
except Exception:
|
|
V = getattr(kv, "vectors", None)
|
|
if V is None:
|
|
raise RuntimeError("KeyedVectors has no 'vectors' matrix.")
|
|
norms = np.linalg.norm(V, axis=1, keepdims=True)
|
|
norms[norms == 0.0] = 1.0
|
|
M = V / norms
|
|
mean_vec = M.mean(axis=0)
|
|
n = np.linalg.norm(mean_vec)
|
|
mean_vec = mean_vec / n if n > 0 else mean_vec
|
|
setattr(kv, "_we_mean_vec", mean_vec)
|
|
return mean_vec
|
|
|
|
def _get_or_make_max_count(kv):
|
|
maxc = getattr(kv, "_we_max_count", None)
|
|
if maxc is not None:
|
|
return maxc
|
|
max_count = 0
|
|
if hasattr(kv, "key_to_index"):
|
|
for w in kv.key_to_index.keys():
|
|
try:
|
|
c = kv.get_vecattr(w, "count")
|
|
if c is not None and c > max_count:
|
|
max_count = int(c)
|
|
except Exception:
|
|
continue
|
|
setattr(kv, "_we_max_count", int(max_count))
|
|
return int(max_count)
|
|
|
|
# ---------- Nodes ----------
|
|
class WordEmbeddingsLoader:
|
|
"""Pretrained (gensim downloader) loader that outputs a WE_MODEL handle."""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
small_default = "glove-wiki-gigaword-50"
|
|
default_choice = (
|
|
small_default if small_default in _MODEL_CHOICES
|
|
else (_MODEL_CHOICES[0] if _MODEL_CHOICES else small_default)
|
|
)
|
|
return {"required": {"model_name": (_MODEL_CHOICES, {"default": default_choice})}}
|
|
|
|
# Single output => easier quick-create suggestions
|
|
RETURN_TYPES = ("WE_MODEL",)
|
|
RETURN_NAMES = ("we_model",)
|
|
FUNCTION = "load"
|
|
CATEGORY = "WordEmbeddings"
|
|
|
|
def load(self, model_name):
|
|
kv = _load_pretrained(model_name)
|
|
return (WEModel(model_name, kv),)
|
|
|
|
class WordEmbeddingsLocalModelLoader:
|
|
"""Local file loader that outputs a WE_MODEL handle."""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {"path": ("STRING", {"default": "", "multiline": False})},
|
|
"optional": {
|
|
"format": (["auto", "kv", "word2vec_bin", "word2vec_text"], {"default": "auto"}),
|
|
"label": ("STRING", {"default": "local", "multiline": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("WE_MODEL",)
|
|
RETURN_NAMES = ("we_model",)
|
|
FUNCTION = "load"
|
|
CATEGORY = "WordEmbeddings"
|
|
|
|
def load(self, path, format="auto", label="local"):
|
|
kv = _load_local_keyedvectors(path, fmt=format)
|
|
name = f"{label}:{os.path.basename(path)}"
|
|
return (WEModel(name, kv),)
|
|
|
|
class WordEmbeddingsExplorer:
|
|
"""Neighbors of a single word from WE_MODEL."""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {"we_model": ("WE_MODEL",)},
|
|
"optional": {
|
|
"word": ("STRING", {"multiline": False, "default": "king"}),
|
|
"k": ("INT", {"default": 10, "min": 1, "max": 50}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING")
|
|
RETURN_NAMES = ("words_csv", "debug_csv")
|
|
FUNCTION = "explore"
|
|
CATEGORY = "WordEmbeddings"
|
|
|
|
def explore(self, we_model, word="king", k=10):
|
|
kv = we_model.kv
|
|
try:
|
|
_ = _kv_get_vec_case_insensitive(kv, word)
|
|
except KeyError:
|
|
return ("", f"Not in vocab: {word}")
|
|
sims = kv.most_similar(word, topn=int(k))
|
|
words = [w for w, _ in sims]
|
|
debug = [f"{w}:{round(score,3)}" for w, score in sims]
|
|
return (",".join(words), ",".join(debug))
|
|
|
|
class WordEmbeddingsInterpolator:
|
|
"""Interpolate between two words on the unit sphere and pick nearest words."""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {"we_model": ("WE_MODEL",)},
|
|
"optional": {
|
|
"word_a": ("STRING", {"multiline": False, "default": "king"}),
|
|
"word_b": ("STRING", {"multiline": False, "default": "queen"}),
|
|
"stops": ("INT", {"default": 12, "min": 1, "max": 256}),
|
|
"topn": ("INT", {"default": 1000, "min": 50, "max": 5000}),
|
|
"method": (["slerp", "lerp"], {"default": "slerp"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING")
|
|
RETURN_NAMES = ("words_csv", "debug_txt")
|
|
FUNCTION = "interpolate"
|
|
CATEGORY = "WordEmbeddings"
|
|
|
|
def interpolate(self, we_model, word_a="king", word_b="queen", stops=12, topn=1000, method="slerp"):
|
|
kv = we_model.kv
|
|
try:
|
|
A = _kv_get_vec_case_insensitive(kv, word_a)
|
|
except KeyError:
|
|
return ("", f"Not in vocab: {word_a}")
|
|
try:
|
|
B = _kv_get_vec_case_insensitive(kv, word_b)
|
|
except KeyError:
|
|
return ("", f"Not in vocab: {word_b}")
|
|
|
|
ts = np.linspace(1.0 / (stops + 1), 1.0 - 1.0 / (stops + 1), int(stops))
|
|
seen = {word_a, word_b, word_a.lower(), word_b.lower()}
|
|
words = []
|
|
debug_lines = []
|
|
|
|
mid = (A + B) / 2
|
|
cos_AB = _cos(A, B)
|
|
|
|
# Endpoint A
|
|
debug_lines.append(
|
|
f"00 t=0.000 word={word_a} cos(v,word)=1.000 cosA=1.000 cosB={cos_AB:.3f} centr={_cos(mid,A):.3f} (endpoint)"
|
|
)
|
|
|
|
# Interior samples
|
|
for i, t in enumerate(ts, 1):
|
|
v = _slerp(A, B, float(t)) if method == "slerp" else ((1 - t) * A + t * B)
|
|
pick = _nearest_to(v, kv, topn=int(topn), seen=seen)
|
|
if not pick:
|
|
debug_lines.append(f"{i:02d} t={t:.3f} (no candidate)")
|
|
continue
|
|
w, x, sim_v = pick
|
|
seen.add(w)
|
|
words.append(w)
|
|
cA = _cos(A, x)
|
|
cB = _cos(B, x)
|
|
cent = _cos(mid, x)
|
|
debug_lines.append(
|
|
f"{i:02d} t={t:.3f} word={w} cos(v,word)={sim_v:.3f} cosA={cA:.3f} cosB={cB:.3f} centr={cent:.3f}"
|
|
)
|
|
|
|
# Endpoint B
|
|
idx_B = len(ts) + 1
|
|
debug_lines.append(
|
|
f"{idx_B:02d} t=1.000 word={word_b} cos(v,word)=1.000 cosA={cos_AB:.3f} cosB=1.000 centr={_cos(mid,B):.3f} (endpoint)"
|
|
)
|
|
|
|
return (",".join(words), "\n".join(debug_lines) if debug_lines else "(no samples)")
|
|
|
|
class WordEmbeddingsEquation:
|
|
"""Vector arithmetic from a free-text expression (e.g., 'king - man + woman')."""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {"we_model": ("WE_MODEL",)},
|
|
"optional": {
|
|
"expression": ("STRING", {"default": "king - man + woman", "multiline": False}),
|
|
"k": ("INT", {"default": 10, "min": 1, "max": 100}),
|
|
"topn_pool": ("INT", {"default": 2000, "min": 100, "max": 20000}),
|
|
"normalize_terms": ("BOOLEAN", {"default": True}),
|
|
"allow_inputs": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING")
|
|
RETURN_NAMES = ("words_csv", "debug_txt")
|
|
FUNCTION = "run"
|
|
CATEGORY = "WordEmbeddings"
|
|
|
|
def _norm(self, v):
|
|
n = np.linalg.norm(v)
|
|
return v / n if n > 0 else v
|
|
|
|
def _parse_expression(self, expr: str):
|
|
if not expr:
|
|
return [], []
|
|
cut = re.split(r"(?:=|equals)", expr, flags=re.IGNORECASE)
|
|
expr = cut[0]
|
|
s = expr.replace(",", " ").replace("\t", " ").strip()
|
|
tokens = [t for t in re.split(r"\s+", s) if t]
|
|
sign = +1.0
|
|
pairs = []
|
|
skipped = []
|
|
i = 0
|
|
while i < len(tokens):
|
|
tok = tokens[i]
|
|
if tok in {"+", "plus"}:
|
|
sign = +1.0; i += 1; continue
|
|
if tok in {"-", "minus"}:
|
|
sign = -1.0; i += 1; continue
|
|
m = re.match(r"^([+-]?\d+(?:\.\d+)?)\*(.+)$", tok)
|
|
if m:
|
|
coef = float(m.group(1)); word = m.group(2)
|
|
pairs.append((word, sign * coef)); sign = +1.0; i += 1; continue
|
|
mnum = re.match(r"^[+-]?\d+(?:\.\d+)?$", tok)
|
|
if mnum and i + 1 < len(tokens):
|
|
try:
|
|
coef = float(tok); word = tokens[i + 1]
|
|
pairs.append((word, sign * coef)); sign = +1.0; i += 2; continue
|
|
except Exception:
|
|
pass
|
|
if tok in {"?", "->", "=>"}:
|
|
i += 1; continue
|
|
pairs.append((tok, sign * 1.0)); sign = +1.0; i += 1
|
|
return pairs, skipped
|
|
|
|
def run(self, we_model, expression="king - man + woman", k=10, topn_pool=2000, normalize_terms=True, allow_inputs=False):
|
|
kv = we_model.kv
|
|
|
|
pairs, skipped = self._parse_expression(expression)
|
|
if not pairs:
|
|
return ("", f"Empty or unparsable expression: '{expression}'")
|
|
|
|
used = []
|
|
missing = []
|
|
vec = None
|
|
|
|
for word, w in pairs:
|
|
try:
|
|
v = _kv_get_vec_case_insensitive(kv, word)
|
|
if normalize_terms:
|
|
v = self._norm(v)
|
|
v = w * v
|
|
vec = v if vec is None else (vec + v)
|
|
used.append((word, w))
|
|
except KeyError:
|
|
missing.append(word)
|
|
|
|
if vec is None:
|
|
miss = ", ".join(missing) if missing else "(none)"
|
|
return ("", f"No valid words found in vocab. Missing: {miss}")
|
|
|
|
target = self._norm(vec)
|
|
|
|
try:
|
|
cand = kv.similar_by_vector(target, topn=int(topn_pool))
|
|
except Exception as e:
|
|
return ("", f"Search error: {e}")
|
|
|
|
input_words = {w for (w, _) in used}
|
|
out = []
|
|
for cw, approx_cos in cand:
|
|
if not allow_inputs and cw in input_words:
|
|
continue
|
|
try:
|
|
x = kv[cw]
|
|
except KeyError:
|
|
continue
|
|
cos_t = float(_cos(target, x))
|
|
out.append((cw, float(approx_cos), cos_t))
|
|
if len(out) >= int(k):
|
|
break
|
|
|
|
if not out:
|
|
dbg = [
|
|
"No results after filtering.",
|
|
f"expr={expression}",
|
|
f"used_terms={used}",
|
|
f"missing_terms={missing}",
|
|
f"allow_inputs={allow_inputs} topn_pool={topn_pool} normalize_terms={normalize_terms}",
|
|
]
|
|
return ("", "\n".join(dbg))
|
|
|
|
words_csv = ",".join([w for (w, _, __) in out])
|
|
debug_lines = [
|
|
f"expr={expression}",
|
|
f"used_terms={used}",
|
|
f"missing_terms={missing}",
|
|
f"normalize_terms={normalize_terms} allow_inputs={allow_inputs} topn_pool={topn_pool}",
|
|
]
|
|
for i, (w2, ac, ct) in enumerate(out, 1):
|
|
debug_lines.append(f"{i:02d} {w2} approx_cos(target,cand)={ac:.3f} cos(target,cand)={ct:.3f}")
|
|
return (words_csv, "\n".join(debug_lines))
|
|
|
|
# -------------------- NEW TOKEN NODES --------------------
|
|
|
|
class WordEmbeddingsTokenNeighbors:
|
|
"""Top-k neighbors for a single token. Outputs words CSV and cosine CSV."""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {"we_model": ("WE_MODEL",)},
|
|
"optional": {
|
|
"word": ("STRING", {"default": "king", "multiline": False}),
|
|
"k": ("INT", {"default": 10, "min": 1, "max": 100}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING")
|
|
RETURN_NAMES = ("words_csv", "cosines_csv")
|
|
FUNCTION = "run"
|
|
CATEGORY = "WordEmbeddings"
|
|
|
|
def run(self, we_model, word="king", k=10):
|
|
kv = we_model.kv
|
|
try:
|
|
_ = _kv_get_vec_case_insensitive(kv, word)
|
|
except KeyError:
|
|
return ("", "")
|
|
sims = kv.most_similar(word, topn=int(k))
|
|
words = [w for w, _ in sims]
|
|
cosines = [f"{float(c):.6f}" for _, c in sims]
|
|
return (",".join(words), ",".join(cosines))
|
|
|
|
# -------- Token Axis (1D) with human summary --------
|
|
class WordEmbeddingsTokenAxis:
|
|
"""
|
|
Project a token onto one semantic axis "left|right".
|
|
Outputs:
|
|
- x : FLOAT in [-1, 1] (positive => toward LEFT pole)
|
|
- summary : STRING (human description combining x with cosine-based confidence)
|
|
- cos_left : FLOAT = cos(token, left_mean)
|
|
- cos_right : FLOAT = cos(token, right_mean)
|
|
|
|
Axis format examples:
|
|
"man,boy,he,him|woman,girl,she,her"
|
|
"royal,king,kingly|common,peasant"
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {"we_model": ("WE_MODEL",)},
|
|
"optional": {
|
|
"token": ("STRING", {"default": "king", "multiline": False}),
|
|
"axis": ("STRING", {
|
|
"default": "man,boy,he,him,his,king,father,brother,husband,actor"
|
|
"|woman,girl,she,her,hers,queen,mother,sister,wife,actress",
|
|
"multiline": True
|
|
}),
|
|
"token_can_be_phrase": ("BOOLEAN", {"default": True}),
|
|
"lowercase": ("BOOLEAN", {"default": True}),
|
|
# how wide a band counts as “similar/equal” on the axis
|
|
"neutral_eps": ("FLOAT", {"default": 0.08, "min": 0.0, "max": 0.5, "step": 0.01}),
|
|
# how tiny x must be to say “equal” (otherwise we say “similar”)
|
|
"equal_eps": ("FLOAT", {"default": 0.02, "min": 0.0, "max": 0.2, "step": 0.005}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("FLOAT", "STRING", "FLOAT", "FLOAT")
|
|
RETURN_NAMES = ("x", "summary", "cos_left", "cos_right")
|
|
FUNCTION = "run"
|
|
CATEGORY = "WordEmbeddings"
|
|
|
|
# ------------ helpers ------------
|
|
@staticmethod
|
|
def _unit(v):
|
|
n = np.linalg.norm(v)
|
|
return v / n if n > 0 else v
|
|
|
|
@staticmethod
|
|
def _split_axis(axis_str: str):
|
|
s = (axis_str or "").strip()
|
|
if "|" in s:
|
|
left_raw, right_raw = s.split("|", 1)
|
|
else:
|
|
left_raw, right_raw = s, ""
|
|
L = [w.strip() for w in left_raw.split(",") if w.strip()]
|
|
R = [w.strip() for w in right_raw.split(",") if w.strip()]
|
|
return L, R, left_raw.strip(), right_raw.strip()
|
|
|
|
@staticmethod
|
|
def _first_term(label_side: str):
|
|
if not label_side:
|
|
return "?"
|
|
return (label_side.split(",")[0].strip() or "?")
|
|
|
|
@classmethod
|
|
def _avg_unit(cls, kv, terms, lowercase=True):
|
|
vecs = []
|
|
for t in terms:
|
|
w = t.lower() if lowercase else t
|
|
try:
|
|
v = _kv_get_vec_case_insensitive(kv, w)
|
|
vecs.append(cls._unit(v))
|
|
except KeyError:
|
|
pass
|
|
if not vecs:
|
|
return None
|
|
return cls._unit(np.mean(vecs, axis=0))
|
|
|
|
@classmethod
|
|
def _phrase_vec(cls, kv, token, token_can_be_phrase=True, lowercase=True):
|
|
pieces = [token] if not token_can_be_phrase else [w for w in token.replace(",", " ").split() if w]
|
|
vecs = []
|
|
for w in pieces:
|
|
ww = w.lower() if lowercase else w
|
|
try:
|
|
vecs.append(cls._unit(_kv_get_vec_case_insensitive(kv, ww)))
|
|
except KeyError:
|
|
pass
|
|
if not vecs:
|
|
return None
|
|
return cls._unit(np.mean(vecs, axis=0))
|
|
|
|
@staticmethod
|
|
def _confidence_phrase(cos_left, cos_right):
|
|
# magnitudes and gap
|
|
C_L, C_R = abs(cos_left), abs(cos_right)
|
|
C_max, C_min = max(C_L, C_R), min(C_L, C_R)
|
|
gap = abs(C_L - C_R)
|
|
|
|
# thresholds (tune if you like)
|
|
BOTH_MIN = 0.35
|
|
BOTH_STRONG = 0.55
|
|
GAP_CLEAR = 0.18
|
|
MAX_MIN = 0.45
|
|
MAX_STRONG = 0.65
|
|
LOW_MAX = 0.25
|
|
|
|
if C_max < LOW_MAX:
|
|
return "with low confidence"
|
|
|
|
# near both poles (high & similar)
|
|
if (C_min >= BOTH_MIN) and (gap < 0.10):
|
|
level = "strong" if C_min >= BOTH_STRONG else "moderate"
|
|
return f"with {level} confidence"
|
|
|
|
# one pole clearly dominates
|
|
if (gap >= GAP_CLEAR) and (C_max >= MAX_MIN):
|
|
side = "left" if (C_L >= C_R) else "right"
|
|
level = "strong" if (C_max >= MAX_STRONG and gap >= 0.30) else "moderate"
|
|
return f"with {level} confidence in {side}"
|
|
|
|
# otherwise, middling
|
|
return "with moderate confidence"
|
|
|
|
# ------------ main ------------
|
|
def run(self, we_model, token="king",
|
|
axis="man,boy,he,him,his,king,father,brother,husband,actor|woman,girl,she,her,hers,queen,mother,sister,wife,actress",
|
|
token_can_be_phrase=True, lowercase=True, neutral_eps=0.08, equal_eps=0.02):
|
|
|
|
kv = we_model.kv
|
|
|
|
# token vector
|
|
v = self._phrase_vec(kv, token, token_can_be_phrase=token_can_be_phrase, lowercase=lowercase)
|
|
if v is None:
|
|
return (0.0, f"{token} not found in the model vocabulary.", 0.0, 0.0)
|
|
|
|
# poles
|
|
L_terms, R_terms, L_raw, R_raw = self._split_axis(axis)
|
|
p = self._avg_unit(kv, L_terms, lowercase=lowercase) if L_terms else None
|
|
n = self._avg_unit(kv, R_terms, lowercase=lowercase) if R_terms else None
|
|
|
|
# cosines to each pole mean
|
|
cos_left = float(v @ p) if p is not None else 0.0
|
|
cos_right = float(v @ n) if n is not None else 0.0
|
|
|
|
# axis direction (fallback to single-sided axis if one pole missing)
|
|
if p is not None and n is not None:
|
|
axis_dir = self._unit(p - n)
|
|
elif p is not None:
|
|
axis_dir = p
|
|
elif n is not None:
|
|
axis_dir = -n
|
|
else:
|
|
# no usable axis
|
|
L = self._first_term(L_raw) if L_raw else "left"
|
|
R = self._first_term(R_raw) if R_raw else "right"
|
|
summary = f"{token} on axis '{L}|{R}': axis unavailable (no pole terms found)."
|
|
return (0.0, summary, cos_left, cos_right)
|
|
|
|
# projection scalar
|
|
x = float(v @ axis_dir)
|
|
|
|
# human summary: keep left/right order
|
|
L = self._first_term(L_raw) if L_raw else "left"
|
|
R = self._first_term(R_raw) if R_raw else "right"
|
|
|
|
# descriptor from x
|
|
x_abs = abs(x)
|
|
if x_abs <= equal_eps:
|
|
dir_phrase = f"equal {L} and {R}"
|
|
elif x_abs < neutral_eps:
|
|
dir_phrase = f"similar {L} and {R}"
|
|
else:
|
|
dir_phrase = (f"more {L} than {R}") if x >= 0 else (f"less {L} than {R}")
|
|
|
|
# confidence from cosines
|
|
conf_phrase = self._confidence_phrase(cos_left, cos_right)
|
|
|
|
summary = f"{token} is {dir_phrase}, {conf_phrase}."
|
|
|
|
return (x, summary, cos_left, cos_right)
|
|
|
|
|
|
class WordEmbeddingsTokenAxis2D:
|
|
"""
|
|
Project a token on TWO semantic axes and emit:
|
|
x (FLOAT), y (FLOAT), summary (STRING),
|
|
x_cos_left, x_cos_right, y_cos_left, y_cos_right (FLOATs).
|
|
|
|
Summary includes only axes considered "strong" (see thresholds).
|
|
Example:
|
|
king is more male than female, with strong confidence; and more royal than common, with moderate confidence in left
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {"we_model": ("WE_MODEL",)},
|
|
"optional": {
|
|
"token": ("STRING", {"default": "king", "multiline": False}),
|
|
"axis1": ("STRING", {"default": "man,boy,he,him,his,father,brother,husband|woman,girl,she,her,hers,mother,sister,wife", "multiline": True}),
|
|
"axis2": ("STRING", {"default": "royal,king,kingly,noble|common,commoner,peasant,ordinary", "multiline": True}),
|
|
"token_can_be_phrase": ("BOOLEAN", {"default": True}),
|
|
"lowercase": ("BOOLEAN", {"default": True}),
|
|
# direction wording
|
|
"neutral_eps": ("FLOAT", {"default": 0.08, "min": 0.0, "max": 0.5, "step": 0.01}),
|
|
"equal_eps": ("FLOAT", {"default": 0.02, "min": 0.0, "max": 0.2, "step": 0.005}),
|
|
# strength gating for summary (axes must pass all to be included)
|
|
"strong_proj": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"strong_cos": ("FLOAT", {"default": 0.65, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"strong_gap": ("FLOAT", {"default": 0.30, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("FLOAT", "FLOAT", "STRING", "FLOAT", "FLOAT", "FLOAT", "FLOAT")
|
|
RETURN_NAMES = ("x", "y", "summary", "x_cos_left", "x_cos_right", "y_cos_left", "y_cos_right")
|
|
FUNCTION = "run"
|
|
CATEGORY = "WordEmbeddings"
|
|
|
|
# ---------- helpers ----------
|
|
@staticmethod
|
|
def _unit(v):
|
|
n = np.linalg.norm(v)
|
|
return v / n if n > 0 else v
|
|
|
|
@staticmethod
|
|
def _split_axis(axis_str: str):
|
|
s = (axis_str or "").strip()
|
|
if "|" in s:
|
|
left_raw, right_raw = s.split("|", 1)
|
|
else:
|
|
left_raw, right_raw = s, ""
|
|
L = [w.strip() for w in left_raw.split(",") if w.strip()]
|
|
R = [w.strip() for w in right_raw.split(",") if w.strip()]
|
|
return L, R, left_raw.strip(), right_raw.strip()
|
|
|
|
@staticmethod
|
|
def _first_term(label_side: str):
|
|
if not label_side:
|
|
return "?"
|
|
return (label_side.split(",")[0].strip() or "?")
|
|
|
|
@classmethod
|
|
def _avg_unit(cls, kv, terms, lowercase=True):
|
|
vecs = []
|
|
for t in terms:
|
|
w = t.lower() if lowercase else t
|
|
try:
|
|
v = _kv_get_vec_case_insensitive(kv, w)
|
|
vecs.append(cls._unit(v))
|
|
except KeyError:
|
|
pass
|
|
if not vecs:
|
|
return None
|
|
return cls._unit(np.mean(vecs, axis=0))
|
|
|
|
@classmethod
|
|
def _phrase_vec(cls, kv, token, token_can_be_phrase=True, lowercase=True):
|
|
parts = [token] if not token_can_be_phrase else [w for w in token.replace(",", " ").split() if w]
|
|
vecs = []
|
|
for w in parts:
|
|
ww = w.lower() if lowercase else w
|
|
try:
|
|
vecs.append(cls._unit(_kv_get_vec_case_insensitive(kv, ww)))
|
|
except KeyError:
|
|
pass
|
|
if not vecs:
|
|
return None
|
|
return cls._unit(np.mean(vecs, axis=0))
|
|
|
|
@staticmethod
|
|
def _direction_phrase(x, L, R, neutral_eps, equal_eps):
|
|
ax = abs(x)
|
|
if ax <= equal_eps:
|
|
return f"equal {L} and {R}"
|
|
if ax < neutral_eps:
|
|
return f"similar {L} and {R}"
|
|
return (f"more {L} than {R}") if x >= 0 else (f"less {L} than {R}")
|
|
|
|
@staticmethod
|
|
def _confidence_level_and_phrase(cos_left, cos_right):
|
|
C_L, C_R = abs(cos_left), abs(cos_right)
|
|
C_max, C_min = max(C_L, C_R), min(C_L, C_R)
|
|
gap = abs(C_L - C_R)
|
|
|
|
BOTH_MIN = 0.35
|
|
BOTH_STRONG = 0.55
|
|
GAP_CLEAR = 0.18
|
|
MAX_MIN = 0.45
|
|
MAX_STRONG = 0.65
|
|
LOW_MAX = 0.25
|
|
|
|
if C_max < LOW_MAX:
|
|
return "low", "with low confidence"
|
|
|
|
if (C_min >= BOTH_MIN) and (gap < 0.10):
|
|
level = "strong" if C_min >= BOTH_STRONG else "moderate"
|
|
return level, f"with {level} confidence"
|
|
|
|
if (gap >= GAP_CLEAR) and (C_max >= MAX_MIN):
|
|
side = "left" if (C_L >= C_R) else "right"
|
|
level = "strong" if (C_max >= MAX_STRONG and gap >= 0.30) else "moderate"
|
|
return level, f"with {level} confidence in {side}"
|
|
|
|
return "moderate", "with moderate confidence"
|
|
|
|
@classmethod
|
|
def _proj_and_cos(cls, kv, v, left_terms, right_terms, lowercase=True):
|
|
Lm = cls._avg_unit(kv, left_terms, lowercase=lowercase) if left_terms else None
|
|
Rm = cls._avg_unit(kv, right_terms, lowercase=lowercase) if right_terms else None
|
|
|
|
cos_left = float(v @ Lm) if (v is not None and Lm is not None) else 0.0
|
|
cos_right = float(v @ Rm) if (v is not None and Rm is not None) else 0.0
|
|
|
|
if Lm is not None and Rm is not None:
|
|
axis_dir = cls._unit(Lm - Rm)
|
|
elif Lm is not None:
|
|
axis_dir = Lm
|
|
elif Rm is not None:
|
|
axis_dir = -Rm
|
|
else:
|
|
return 0.0, cos_left, cos_right, False
|
|
|
|
proj = float(v @ axis_dir) if v is not None else 0.0
|
|
return proj, cos_left, cos_right, True
|
|
|
|
# ---------- main ----------
|
|
def run(self, we_model, token="king",
|
|
axis1="man,boy,he,him,his,father,brother,husband|woman,girl,she,her,hers,mother,sister,wife",
|
|
axis2="royal,king,kingly,noble|common,commoner,peasant,ordinary",
|
|
token_can_be_phrase=True, lowercase=True,
|
|
neutral_eps=0.08, equal_eps=0.02,
|
|
strong_proj=0.35, strong_cos=0.65, strong_gap=0.30):
|
|
|
|
kv = we_model.kv
|
|
v = self._phrase_vec(kv, token, token_can_be_phrase=token_can_be_phrase, lowercase=lowercase)
|
|
if v is None:
|
|
return (0.0, 0.0, f"{token} not found in the model vocabulary.", 0.0, 0.0, 0.0, 0.0)
|
|
|
|
# parse axes
|
|
a1L, a1R, a1L_raw, a1R_raw = self._split_axis(axis1)
|
|
a2L, a2R, a2L_raw, a2R_raw = self._split_axis(axis2)
|
|
|
|
# compute projections & cosines
|
|
x, x_cos_left, x_cos_right, ok1 = self._proj_and_cos(kv, v, a1L, a1R, lowercase=lowercase)
|
|
y, y_cos_left, y_cos_right, ok2 = self._proj_and_cos(kv, v, a2L, a2R, lowercase=lowercase)
|
|
|
|
# build per-axis clauses (but only include "strong" ones)
|
|
clauses = []
|
|
|
|
if ok1:
|
|
dir1 = self._direction_phrase(x, self._first_term(a1L_raw), self._first_term(a1R_raw), neutral_eps, equal_eps)
|
|
lvl1, conf1 = self._confidence_level_and_phrase(x_cos_left, x_cos_right)
|
|
strong1 = (abs(x) >= strong_proj) and (
|
|
max(abs(x_cos_left), abs(x_cos_right)) >= strong_cos or
|
|
(max(abs(x_cos_left), abs(x_cos_right)) >= 0.55 and abs(abs(x_cos_left) - abs(x_cos_right)) >= strong_gap)
|
|
)
|
|
if strong1:
|
|
clauses.append(f"{dir1}, {conf1}")
|
|
|
|
if ok2:
|
|
dir2 = self._direction_phrase(y, self._first_term(a2L_raw), self._first_term(a2R_raw), neutral_eps, equal_eps)
|
|
lvl2, conf2 = self._confidence_level_and_phrase(y_cos_left, y_cos_right)
|
|
strong2 = (abs(y) >= strong_proj) and (
|
|
max(abs(y_cos_left), abs(y_cos_right)) >= strong_cos or
|
|
(max(abs(y_cos_left), abs(y_cos_right)) >= 0.55 and abs(abs(y_cos_left) - abs(y_cos_right)) >= strong_gap)
|
|
)
|
|
if strong2:
|
|
clauses.append(f"{dir2}, {conf2}")
|
|
|
|
if not clauses:
|
|
summary = f"{token} shows no strong alignment on the provided axes."
|
|
elif len(clauses) == 1:
|
|
summary = f"{token} is {clauses[0]}"
|
|
else:
|
|
summary = f"{token} is {clauses[0]} and {clauses[1]}"
|
|
|
|
return (x, y, summary, x_cos_left, x_cos_right, y_cos_left, y_cos_right)
|
|
|
|
|
|
class WordEmbeddingsTokenAxis3D:
|
|
"""
|
|
Project a token on THREE semantic axes and emit:
|
|
x, y, z (FLOATs), summary (STRING),
|
|
x_cos_left/x_cos_right, y_cos_left/y_cos_right, z_cos_left/z_cos_right (FLOATs).
|
|
|
|
Summary includes only axes considered "strong" (see thresholds).
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {"we_model": ("WE_MODEL",)},
|
|
"optional": {
|
|
"token": ("STRING", {"default": "king", "multiline": False}),
|
|
"axis1": ("STRING", {"default": "man,boy,he,him,his,father,brother,husband|woman,girl,she,her,hers,mother,sister,wife", "multiline": True}),
|
|
"axis2": ("STRING", {"default": "hot,warm,heat|cold,cool,freezing", "multiline": True}),
|
|
"axis3": ("STRING", {"default": "royal,king,kingly,noble|common,commoner,peasant,ordinary", "multiline": True}),
|
|
"token_can_be_phrase": ("BOOLEAN", {"default": True}),
|
|
"lowercase": ("BOOLEAN", {"default": True}),
|
|
# direction wording
|
|
"neutral_eps": ("FLOAT", {"default": 0.08, "min": 0.0, "max": 0.5, "step": 0.01}),
|
|
"equal_eps": ("FLOAT", {"default": 0.02, "min": 0.0, "max": 0.2, "step": 0.005}),
|
|
# strength gating for summary
|
|
"strong_proj": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"strong_cos": ("FLOAT", {"default": 0.65, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"strong_gap": ("FLOAT", {"default": 0.30, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (
|
|
"FLOAT","FLOAT","FLOAT","STRING",
|
|
"FLOAT","FLOAT","FLOAT","FLOAT","FLOAT","FLOAT"
|
|
)
|
|
RETURN_NAMES = (
|
|
"x","y","z","summary",
|
|
"x_cos_left","x_cos_right","y_cos_left","y_cos_right","z_cos_left","z_cos_right"
|
|
)
|
|
FUNCTION = "run"
|
|
CATEGORY = "WordEmbeddings"
|
|
|
|
# ---- reuse helpers from 2D ----
|
|
@staticmethod
|
|
def _unit(v):
|
|
n = np.linalg.norm(v)
|
|
return v / n if n > 0 else v
|
|
|
|
@staticmethod
|
|
def _split_axis(axis_str: str):
|
|
s = (axis_str or "").strip()
|
|
if "|" in s:
|
|
left_raw, right_raw = s.split("|", 1)
|
|
else:
|
|
left_raw, right_raw = s, ""
|
|
L = [w.strip() for w in left_raw.split(",") if w.strip()]
|
|
R = [w.strip() for w in right_raw.split(",") if w.strip()]
|
|
return L, R, left_raw.strip(), right_raw.strip()
|
|
|
|
@staticmethod
|
|
def _first_term(label_side: str):
|
|
if not label_side:
|
|
return "?"
|
|
return (label_side.split(",")[0].strip() or "?")
|
|
|
|
@classmethod
|
|
def _avg_unit(cls, kv, terms, lowercase=True):
|
|
vecs = []
|
|
for t in terms:
|
|
w = t.lower() if lowercase else t
|
|
try:
|
|
v = _kv_get_vec_case_insensitive(kv, w)
|
|
vecs.append(cls._unit(v))
|
|
except KeyError:
|
|
pass
|
|
if not vecs:
|
|
return None
|
|
return cls._unit(np.mean(vecs, axis=0))
|
|
|
|
@classmethod
|
|
def _phrase_vec(cls, kv, token, token_can_be_phrase=True, lowercase=True):
|
|
parts = [token] if not token_can_be_phrase else [w for w in token.replace(",", " ").split() if w]
|
|
vecs = []
|
|
for w in parts:
|
|
ww = w.lower() if lowercase else w
|
|
try:
|
|
vecs.append(cls._unit(_kv_get_vec_case_insensitive(kv, ww)))
|
|
except KeyError:
|
|
pass
|
|
if not vecs:
|
|
return None
|
|
return cls._unit(np.mean(vecs, axis=0))
|
|
|
|
@staticmethod
|
|
def _direction_phrase(x, L, R, neutral_eps, equal_eps):
|
|
ax = abs(x)
|
|
if ax <= equal_eps:
|
|
return f"equal {L} and {R}"
|
|
if ax < neutral_eps:
|
|
return f"similar {L} and {R}"
|
|
return (f"more {L} than {R}") if x >= 0 else (f"less {L} than {R}")
|
|
|
|
@staticmethod
|
|
def _confidence_level_and_phrase(cos_left, cos_right):
|
|
C_L, C_R = abs(cos_left), abs(cos_right)
|
|
C_max, C_min = max(C_L, C_R), min(C_L, C_R)
|
|
gap = abs(C_L - C_R)
|
|
|
|
BOTH_MIN = 0.35
|
|
BOTH_STRONG = 0.55
|
|
GAP_CLEAR = 0.18
|
|
MAX_MIN = 0.45
|
|
MAX_STRONG = 0.65
|
|
LOW_MAX = 0.25
|
|
|
|
if C_max < LOW_MAX:
|
|
return "low", "with low confidence"
|
|
|
|
if (C_min >= BOTH_MIN) and (gap < 0.10):
|
|
level = "strong" if C_min >= BOTH_STRONG else "moderate"
|
|
return level, f"with {level} confidence"
|
|
|
|
if (gap >= GAP_CLEAR) and (C_max >= MAX_MIN):
|
|
side = "left" if (C_L >= C_R) else "right"
|
|
level = "strong" if (C_max >= MAX_STRONG and gap >= 0.30) else "moderate"
|
|
return level, f"with {level} confidence in {side}"
|
|
|
|
return "moderate", "with moderate confidence"
|
|
|
|
@classmethod
|
|
def _proj_and_cos(cls, kv, v, left_terms, right_terms, lowercase=True):
|
|
Lm = cls._avg_unit(kv, left_terms, lowercase=lowercase) if left_terms else None
|
|
Rm = cls._avg_unit(kv, right_terms, lowercase=lowercase) if right_terms else None
|
|
|
|
cos_left = float(v @ Lm) if (v is not None and Lm is not None) else 0.0
|
|
cos_right = float(v @ Rm) if (v is not None and Rm is not None) else 0.0
|
|
|
|
if Lm is not None and Rm is not None:
|
|
axis_dir = cls._unit(Lm - Rm)
|
|
elif Lm is not None:
|
|
axis_dir = Lm
|
|
elif Rm is not None:
|
|
axis_dir = -Rm
|
|
else:
|
|
return 0.0, cos_left, cos_right, False
|
|
|
|
proj = float(v @ axis_dir) if v is not None else 0.0
|
|
return proj, cos_left, cos_right, True
|
|
|
|
def run(self, we_model, token="king",
|
|
axis1="man,boy,he,him,his,father,brother,husband|woman,girl,she,her,hers,mother,sister,wife",
|
|
axis2="hot,warm,heat|cold,cool,freezing",
|
|
axis3="royal,king,kingly,noble|common,commoner,peasant,ordinary",
|
|
token_can_be_phrase=True, lowercase=True,
|
|
neutral_eps=0.08, equal_eps=0.02,
|
|
strong_proj=0.35, strong_cos=0.65, strong_gap=0.30):
|
|
|
|
kv = we_model.kv
|
|
v = self._phrase_vec(kv, token, token_can_be_phrase=token_can_be_phrase, lowercase=lowercase)
|
|
if v is None:
|
|
return (
|
|
0.0, 0.0, 0.0,
|
|
f"{token} not found in the model vocabulary.",
|
|
0.0, 0.0, 0.0, 0.0, 0.0, 0.0
|
|
)
|
|
|
|
# parse
|
|
a1L, a1R, a1L_raw, a1R_raw = self._split_axis(axis1)
|
|
a2L, a2R, a2L_raw, a2R_raw = self._split_axis(axis2)
|
|
a3L, a3R, a3L_raw, a3R_raw = self._split_axis(axis3)
|
|
|
|
# projections & cosines
|
|
x, x_cos_left, x_cos_right, ok1 = self._proj_and_cos(kv, v, a1L, a1R, lowercase=lowercase)
|
|
y, y_cos_left, y_cos_right, ok2 = self._proj_and_cos(kv, v, a2L, a2R, lowercase=lowercase)
|
|
z, z_cos_left, z_cos_right, ok3 = self._proj_and_cos(kv, v, a3L, a3R, lowercase=lowercase)
|
|
|
|
# build clauses only for strong axes
|
|
clauses = []
|
|
|
|
if ok1:
|
|
dir1 = self._direction_phrase(x, self._first_term(a1L_raw), self._first_term(a1R_raw), neutral_eps, equal_eps)
|
|
lvl1, conf1 = self._confidence_level_and_phrase(x_cos_left, x_cos_right)
|
|
strong1 = (abs(x) >= strong_proj) and (
|
|
max(abs(x_cos_left), abs(x_cos_right)) >= strong_cos or
|
|
(max(abs(x_cos_left), abs(x_cos_right)) >= 0.55 and abs(abs(x_cos_left) - abs(x_cos_right)) >= strong_gap)
|
|
)
|
|
if strong1:
|
|
clauses.append(f"{dir1}, {conf1}")
|
|
|
|
if ok2:
|
|
dir2 = self._direction_phrase(y, self._first_term(a2L_raw), self._first_term(a2R_raw), neutral_eps, equal_eps)
|
|
lvl2, conf2 = self._confidence_level_and_phrase(y_cos_left, y_cos_right)
|
|
strong2 = (abs(y) >= strong_proj) and (
|
|
max(abs(y_cos_left), abs(y_cos_right)) >= strong_cos or
|
|
(max(abs(y_cos_left), abs(y_cos_right)) >= 0.55 and abs(abs(y_cos_left) - abs(y_cos_right)) >= strong_gap)
|
|
)
|
|
if strong2:
|
|
clauses.append(f"{dir2}, {conf2}")
|
|
|
|
if ok3:
|
|
dir3 = self._direction_phrase(z, self._first_term(a3L_raw), self._first_term(a3R_raw), neutral_eps, equal_eps)
|
|
lvl3, conf3 = self._confidence_level_and_phrase(z_cos_left, z_cos_right)
|
|
strong3 = (abs(z) >= strong_proj) and (
|
|
max(abs(z_cos_left), abs(z_cos_right)) >= strong_cos or
|
|
(max(abs(z_cos_left), abs(z_cos_right)) >= 0.55 and abs(abs(z_cos_left) - abs(z_cos_right)) >= strong_gap)
|
|
)
|
|
if strong3:
|
|
clauses.append(f"{dir3}, {conf3}")
|
|
|
|
if not clauses:
|
|
summary = f"{token} shows no strong alignment on the provided axes."
|
|
elif len(clauses) == 1:
|
|
summary = f"{token} is {clauses[0]}"
|
|
elif len(clauses) == 2:
|
|
summary = f"{token} is {clauses[0]} and {clauses[1]}"
|
|
else:
|
|
summary = f"{token} is {clauses[0]}, {clauses[1]}, and {clauses[2]}"
|
|
|
|
return (
|
|
x, y, z, summary,
|
|
x_cos_left, x_cos_right, y_cos_left, y_cos_right, z_cos_left, z_cos_right
|
|
)
|
|
|
|
|
|
class WordEmbeddingsTokenCentrality:
|
|
"""
|
|
Centrality of a token relative to the model's global mean direction.
|
|
|
|
Outputs:
|
|
- centrality (FLOAT in [-1, 1]) : cosine(v̂, mean̂)
|
|
- norm (FLOAT) : vector length of the raw token embedding
|
|
- summary (STRING) : human-friendly description
|
|
|
|
Notes:
|
|
- |centrality| ≈ how "typical/generic" the token is in this space.
|
|
* high |c| -> very central / typical
|
|
* mid |c| -> somewhat central
|
|
* low |c| -> distinctive / outlier
|
|
- The sign of centrality is usually not semantically meaningful across models,
|
|
but you can optionally mention orientation if desired.
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {"we_model": ("WE_MODEL",)},
|
|
"optional": {
|
|
"token": ("STRING", {"default": "king", "multiline": False}),
|
|
"token_can_be_phrase": ("BOOLEAN", {"default": True}),
|
|
"lowercase": ("BOOLEAN", {"default": True}),
|
|
# Summary tuning
|
|
"strong_thr": ("FLOAT", {"default": 0.65, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"moderate_thr": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"include_orientation": ("BOOLEAN", {"default": False}), # mention sign of centrality
|
|
"orientation_eps": ("FLOAT", {"default": 0.05, "min": 0.0, "max": 0.2, "step": 0.005}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("FLOAT", "FLOAT", "STRING")
|
|
RETURN_NAMES = ("centrality", "norm", "summary")
|
|
FUNCTION = "run"
|
|
CATEGORY = "WordEmbeddings"
|
|
|
|
# ---------- helpers ----------
|
|
@staticmethod
|
|
def _unit(v):
|
|
n = np.linalg.norm(v)
|
|
return v / n if n > 0 else v
|
|
|
|
@classmethod
|
|
def _phrase_vec(cls, kv, token, token_can_be_phrase=True, lowercase=True):
|
|
parts = [token] if not token_can_be_phrase else [w for w in token.replace(",", " ").split() if w]
|
|
vecs = []
|
|
for w in parts:
|
|
ww = w.lower() if lowercase else w
|
|
try:
|
|
vecs.append(cls._unit(_kv_get_vec_case_insensitive(kv, ww)))
|
|
except KeyError:
|
|
pass
|
|
if not vecs:
|
|
return None, 0.0
|
|
raw = np.mean(vecs, axis=0)
|
|
return cls._unit(raw), float(np.linalg.norm(raw))
|
|
|
|
@staticmethod
|
|
def _tier_phrase(mag, strong_thr, moderate_thr):
|
|
if mag >= strong_thr:
|
|
return "highly central (very typical)"
|
|
if mag >= moderate_thr:
|
|
return "moderately central"
|
|
return "weakly central (distinctive/outlier)"
|
|
|
|
@staticmethod
|
|
def _orientation_phrase(cent, orientation_eps):
|
|
if cent >= orientation_eps:
|
|
return "with the average orientation"
|
|
if cent <= -orientation_eps:
|
|
return "with opposite-to-average orientation"
|
|
return "with near-orthogonal orientation"
|
|
|
|
# ---------- main ----------
|
|
def run(self, we_model,
|
|
token="king",
|
|
token_can_be_phrase=True, lowercase=True,
|
|
strong_thr=0.65, moderate_thr=0.35,
|
|
include_orientation=False, orientation_eps=0.05):
|
|
|
|
kv = we_model.kv
|
|
|
|
vhat, norm_val = self._phrase_vec(kv, token, token_can_be_phrase=token_can_be_phrase, lowercase=lowercase)
|
|
if vhat is None:
|
|
return (0.0, 0.0, f"{token} not found in the model vocabulary.")
|
|
|
|
mv = _get_or_make_mean_vec(kv) # already unit-length
|
|
cent = float(vhat @ mv) # cosine in [-1, 1]
|
|
mag = abs(cent)
|
|
|
|
tier = self._tier_phrase(mag, strong_thr, moderate_thr)
|
|
if include_orientation:
|
|
orient = self._orientation_phrase(cent, orientation_eps)
|
|
summary = f"{token} is {tier}, {orient}"
|
|
else:
|
|
summary = f"{token} is {tier}"
|
|
|
|
return (cent, norm_val, summary)
|
|
|
|
|
|
# ---------- Registration ----------
|
|
NODE_CLASS_MAPPINGS = {
|
|
"WordEmbeddingsLoader": WordEmbeddingsLoader,
|
|
"WordEmbeddingsLocalModelLoader": WordEmbeddingsLocalModelLoader,
|
|
"WordEmbeddingsExplorer": WordEmbeddingsExplorer,
|
|
"WordEmbeddingsInterpolator": WordEmbeddingsInterpolator,
|
|
"WordEmbeddingsEquation": WordEmbeddingsEquation,
|
|
|
|
# new token nodes
|
|
"WordEmbeddingsTokenNeighbors": WordEmbeddingsTokenNeighbors,
|
|
"WordEmbeddingsTokenAxis": WordEmbeddingsTokenAxis,
|
|
"WordEmbeddingsTokenAxis2D": WordEmbeddingsTokenAxis2D,
|
|
"WordEmbeddingsTokenAxis3D": WordEmbeddingsTokenAxis3D,
|
|
"WordEmbeddingsTokenCentrality": WordEmbeddingsTokenCentrality,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"WordEmbeddingsLoader": "WordEmbeddings: Loader",
|
|
"WordEmbeddingsLocalModelLoader": "WordEmbeddings: Local Loader",
|
|
"WordEmbeddingsExplorer": "WordEmbeddings: Explorer",
|
|
"WordEmbeddingsInterpolator": "WordEmbeddings: Interpolator",
|
|
"WordEmbeddingsEquation": "WordEmbeddings: Equation",
|
|
|
|
# new token nodes
|
|
"WordEmbeddingsTokenNeighbors": "WordEmbeddings: Token Neighbors",
|
|
"WordEmbeddingsTokenAxis": "WordEmbeddings: Token Axis",
|
|
"WordEmbeddingsTokenAxis2D": "WordEmbeddings: Axis 2D",
|
|
"WordEmbeddingsTokenAxis3D": "WordEmbeddings: Axis 3D",
|
|
"WordEmbeddingsTokenCentrality": "WordEmbeddings: Token Centrality",
|
|
}
|