Updated refiner logic + removed obsolete functionality
This commit is contained in:
aimfordeb
+135
-423
@@ -43,6 +43,7 @@ def serializable_to_tensor(d: dict) -> torch.Tensor:
|
||||
if torch.cuda.is_available():
|
||||
tensor = tensor.cuda()
|
||||
return tensor
|
||||
|
||||
class FunPackGemmaEmbeddingRefiner:
|
||||
_tokenizer = None
|
||||
|
||||
@@ -82,13 +83,6 @@ class FunPackGemmaEmbeddingRefiner:
|
||||
"default": "",
|
||||
"placeholder": "Positive prompt"
|
||||
}),
|
||||
"negative_prompt": ("STRING", {
|
||||
"multiline": True,
|
||||
"default": "",
|
||||
"placeholder": "Negative prompt (optional)"
|
||||
}),
|
||||
"negative_conditioning": ("CONDITIONING",),
|
||||
"latent": ("LATENT",),
|
||||
"reset_session": ("BOOLEAN", {"default": False, "label": "Reset Session (clears ALL history)"}),
|
||||
"unlimited_history": ("BOOLEAN", {
|
||||
"default": False,
|
||||
@@ -107,8 +101,8 @@ class FunPackGemmaEmbeddingRefiner:
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT", "STRING", "STRING", "STRING")
|
||||
RETURN_NAMES = ("modified_positive", "modified_negative", "modified_latent", "status", "feedback_question", "training_info")
|
||||
RETURN_TYPES = ("CONDITIONING", "STRING", "STRING", "STRING")
|
||||
RETURN_NAMES = ("modified_positive", "status", "feedback_question", "training_info")
|
||||
FUNCTION = "refine"
|
||||
CATEGORY = "FunPack/Refinement"
|
||||
|
||||
@@ -117,7 +111,7 @@ class FunPackGemmaEmbeddingRefiner:
|
||||
return float("nan")
|
||||
|
||||
def _get_scheduler_factors(self, mode, rating, reward, similarity, iter_num, total_iters,
|
||||
token_importance, word_groups, full_token_ids, global_adaptive, device):
|
||||
word_importance, word_groups, global_adaptive, device):
|
||||
if mode == "original":
|
||||
base_lr = 0.18 / (1 + 0.08 * (iter_num ** 0.5))
|
||||
confidence = max(0.2, min(1.0, (rating - 3.0) / 5.0))
|
||||
@@ -133,17 +127,17 @@ class FunPackGemmaEmbeddingRefiner:
|
||||
d_coef = 1.0 if mode == "accurate" else 1.8
|
||||
adaptive_lr = global_adaptive.get("prodigy_lr_base", 1.0)
|
||||
|
||||
token_lr_mult = {}
|
||||
for tid in full_token_ids:
|
||||
tid_str = str(tid)
|
||||
if tid_str not in token_importance:
|
||||
word_lr_mult = {}
|
||||
for _, _, full_word, _ in word_groups:
|
||||
wkey = full_word.lower()
|
||||
if wkey not in word_importance:
|
||||
continue
|
||||
g = abs(reward) + 1e-8
|
||||
if tid_str not in prodigy_d:
|
||||
prodigy_d[tid_str] = g
|
||||
if wkey not in prodigy_d:
|
||||
prodigy_d[wkey] = g
|
||||
else:
|
||||
prodigy_d[tid_str] = 0.9 * prodigy_d[tid_str] + 0.1 * g
|
||||
token_lr_mult[tid_str] = adaptive_lr / (prodigy_d[tid_str] ** 0.5 + 1e-8) * d_coef
|
||||
prodigy_d[wkey] = 0.9 * prodigy_d[wkey] + 0.1 * g
|
||||
word_lr_mult[wkey] = adaptive_lr / (prodigy_d[wkey] ** 0.5 + 1e-8) * d_coef
|
||||
|
||||
current_step = global_adaptive.setdefault("current_step", 0)
|
||||
current_step = min(current_step + 1, 500)
|
||||
@@ -166,20 +160,20 @@ class FunPackGemmaEmbeddingRefiner:
|
||||
confidence = min(1.0, confidence * 1.3)
|
||||
exploration_mult = max(0.6, 1.3 - 0.7 * progress)
|
||||
|
||||
return lr_scale, confidence, exploration_mult, token_lr_mult
|
||||
return lr_scale, confidence, exploration_mult, word_lr_mult
|
||||
|
||||
def _get_top_tokens(self, token_dict, tokenizer, top_k=10):
|
||||
if not token_dict or not tokenizer:
|
||||
return "N/A"
|
||||
sorted_tokens = sorted(token_dict.items(), key=lambda x: x[1], reverse=True)[:top_k]
|
||||
top_list = []
|
||||
for tid_str, score in sorted_tokens:
|
||||
for key, score in sorted_tokens:
|
||||
try:
|
||||
token_text = tokenizer.decode([int(tid_str)], skip_special_tokens=True).strip()
|
||||
if token_text:
|
||||
top_list.append(f"{token_text}({score:.2f})")
|
||||
text = key if isinstance(key, str) else tokenizer.decode([int(key)], skip_special_tokens=True).strip()
|
||||
if text:
|
||||
top_list.append(f"{text}({score:.2f})")
|
||||
except:
|
||||
top_list.append(f"ID{tid_str}({score:.2f})")
|
||||
top_list.append(f"{key}({score:.2f})")
|
||||
return ", ".join(top_list) if top_list else "None"
|
||||
|
||||
def _is_valuable_token(self, token_text):
|
||||
@@ -199,8 +193,7 @@ class FunPackGemmaEmbeddingRefiner:
|
||||
return True
|
||||
|
||||
def refine(self, positive_conditioning, rating: int, refinement_key: str, scheduler_mode: str = "original",
|
||||
positive_prompt: str = "", negative_prompt: str = "",
|
||||
negative_conditioning=None, latent=None,
|
||||
positive_prompt: str = "",
|
||||
reset_session: bool = False, unlimited_history: bool = False,
|
||||
seed: int = 0,
|
||||
feedback_enabled: bool = False, feedback_rating: int = 3):
|
||||
@@ -217,7 +210,7 @@ class FunPackGemmaEmbeddingRefiner:
|
||||
json_file = os.path.join(refinements_dir, f"refine_{safe_key}.json")
|
||||
|
||||
if not positive_conditioning or not isinstance(positive_conditioning, list) or len(positive_conditioning) == 0:
|
||||
return (positive_conditioning, negative_conditioning, latent, "ERROR: Empty positive CONDITIONING input", "", "ERROR: No positive conditioning")
|
||||
return (positive_conditioning, "ERROR: Empty positive CONDITIONING input", "", "ERROR: No positive conditioning")
|
||||
|
||||
item = positive_conditioning[0]
|
||||
if isinstance(item, (list, tuple)) and len(item) >= 2:
|
||||
@@ -228,32 +221,44 @@ class FunPackGemmaEmbeddingRefiner:
|
||||
positive_meta = {"pooled_output": None}
|
||||
|
||||
if not isinstance(raw_positive, torch.Tensor):
|
||||
return (positive_conditioning, negative_conditioning, latent, "ERROR: No positive embedding tensor found", "", "ERROR: Invalid embedding")
|
||||
return (positive_conditioning, "ERROR: No positive embedding tensor found", "", "ERROR: Invalid embedding")
|
||||
|
||||
prompt_key = positive_prompt
|
||||
if negative_prompt:
|
||||
prompt_key += f" |||NEG||| {negative_prompt}"
|
||||
|
||||
if reset_session or not os.path.exists(json_file):
|
||||
data = {
|
||||
"refinement_key": refinement_key,
|
||||
"global_adaptive": {
|
||||
"token_importance": {}, "infusion_tokens": {}, "token_library": {},
|
||||
"exploration_base": 0.08, "momentum": None, "avg_reward_ema": 0.0,
|
||||
"good_ratio": 0.0, "dynamic_sim_threshold": 0.82,
|
||||
"last_feedback_tid": None, "last_feedback_word": None, "feedbacked_words": [],
|
||||
"scheduler_mode": scheduler_mode, "prodigy_d": {}, "prodigy_lr_base": 1.0,
|
||||
"warmup_steps": 8, "total_steps_estimate": 150, "current_step": 0,
|
||||
"word_importance": {},
|
||||
"infusion_tokens": {},
|
||||
"token_library": {},
|
||||
"exploration_base": 0.08,
|
||||
"momentum": None,
|
||||
"avg_reward_ema": 0.0,
|
||||
"good_ratio": 0.0,
|
||||
"dynamic_sim_threshold": 0.82,
|
||||
"last_feedback_word": None,
|
||||
"feedbacked_words": [],
|
||||
"scheduler_mode": scheduler_mode,
|
||||
"prodigy_d": {},
|
||||
"prodigy_lr_base": 1.0,
|
||||
"warmup_steps": 8,
|
||||
"total_steps_estimate": 150,
|
||||
"current_step": 0,
|
||||
},
|
||||
"prompt_histories": {
|
||||
prompt_key: {"reference_embeds": tensor_to_serializable(raw_positive), "history": [], "last_rating": rating}
|
||||
prompt_key: {
|
||||
"reference_embeds": tensor_to_serializable(raw_positive),
|
||||
"history": [],
|
||||
"last_rating": rating
|
||||
}
|
||||
},
|
||||
"last_prompt_key": prompt_key,
|
||||
"pending_feedback": None
|
||||
}
|
||||
with open(json_file, "w", encoding="utf-8") as f:
|
||||
json.dump(data, f, indent=2)
|
||||
return (positive_conditioning, negative_conditioning, latent, "✓ New session started – Reference saved", "", "✓ New session started. Reference embedding saved.")
|
||||
return (positive_conditioning, "✓ New session started – Reference saved", "", "✓ New session started. Reference embedding saved.")
|
||||
|
||||
with open(json_file, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
@@ -262,7 +267,7 @@ class FunPackGemmaEmbeddingRefiner:
|
||||
prompt_histories = data.get("prompt_histories", {})
|
||||
tokenizer = self._get_tokenizer()
|
||||
|
||||
# ====================== FEEDBACK STATE MACHINE ======================
|
||||
# ====================== FEEDBACK STATE MACHINE (WORD-LEVEL) ======================
|
||||
feedback_question_output = ""
|
||||
pending = data.get("pending_feedback")
|
||||
if not feedback_enabled:
|
||||
@@ -273,23 +278,16 @@ class FunPackGemmaEmbeddingRefiner:
|
||||
feedback_question_output = "Feedback disabled. Queue cleared."
|
||||
else:
|
||||
if pending is not None:
|
||||
token_id = pending["token_id"]
|
||||
tid_str = str(token_id)
|
||||
token_importance = global_adaptive["token_importance"]
|
||||
old_imp = token_importance.get(tid_str, 1.0)
|
||||
full_word = pending["full_word"].lower()
|
||||
word_importance = global_adaptive["word_importance"]
|
||||
|
||||
old_imp = word_importance.get(full_word, 1.0)
|
||||
if feedback_rating == 6:
|
||||
target_imp = 0.65
|
||||
else:
|
||||
target_imp = 0.4 + (feedback_rating / 5.0) * 2.1
|
||||
new_imp = 0.65 * old_imp + 0.35 * target_imp
|
||||
token_importance[tid_str] = max(0.3, min(2.5, new_imp))
|
||||
|
||||
token_library = global_adaptive.setdefault("token_library", {})
|
||||
if tid_str in token_library:
|
||||
if feedback_rating >= 4 and feedback_rating != 6:
|
||||
token_library[tid_str]["high_rated_count"] += 2
|
||||
elif feedback_rating <= 2:
|
||||
token_library[tid_str]["low_rated_count"] += 2
|
||||
word_importance[full_word] = max(0.3, min(2.8, new_imp))
|
||||
|
||||
data["pending_feedback"] = None
|
||||
with open(json_file, "w", encoding="utf-8") as f:
|
||||
@@ -303,7 +301,11 @@ class FunPackGemmaEmbeddingRefiner:
|
||||
is_new_prompt = False
|
||||
else:
|
||||
is_new_prompt = True
|
||||
active = {"reference_embeds": tensor_to_serializable(raw_positive), "history": [], "last_rating": rating}
|
||||
active = {
|
||||
"reference_embeds": tensor_to_serializable(raw_positive),
|
||||
"history": [],
|
||||
"last_rating": rating
|
||||
}
|
||||
prompt_histories[prompt_key] = active
|
||||
|
||||
old_reference = serializable_to_tensor(active["reference_embeds"])
|
||||
@@ -312,71 +314,28 @@ class FunPackGemmaEmbeddingRefiner:
|
||||
|
||||
full_token_ids = tokenizer.encode(positive_prompt, add_special_tokens=True) if tokenizer and positive_prompt else []
|
||||
|
||||
# ====================== IMPROVED WORD GROUPING (FIXED) ======================
|
||||
word_groups = []
|
||||
if tokenizer and positive_prompt and full_token_ids:
|
||||
try:
|
||||
# Get raw tokens with Gemma's ▁ word-boundary marker
|
||||
token_texts = tokenizer.convert_ids_to_tokens(full_token_ids)
|
||||
i = 0
|
||||
while i < len(token_texts):
|
||||
start_idx = i
|
||||
word_token_ids = [full_token_ids[i]]
|
||||
# Build the human-readable word
|
||||
word_str = token_texts[i].replace("▁", "")
|
||||
|
||||
i += 1
|
||||
while i < len(token_texts) and not token_texts[i].startswith("▁"):
|
||||
word_token_ids.append(full_token_ids[i])
|
||||
word_str += token_texts[i].replace("▁", "")
|
||||
i += 1
|
||||
|
||||
# Final clean decoded word for feedback
|
||||
full_word = tokenizer.decode(word_token_ids, skip_special_tokens=True).strip()
|
||||
if not full_word:
|
||||
full_word = word_str.strip()
|
||||
|
||||
word_groups.append((start_idx, i, full_word, word_token_ids))
|
||||
except Exception:
|
||||
# Ultra-safe fallback
|
||||
word_groups = [(i, i+1, tokenizer.decode([tid]).strip(), [tid]) for i, tid in enumerate(full_token_ids)]
|
||||
|
||||
# ====================== CROSS-PROMPT TRANSFER ======================
|
||||
if is_new_prompt and tokenizer and full_token_ids:
|
||||
current_set = set(full_token_ids)
|
||||
token_importance = global_adaptive["token_importance"]
|
||||
for other_key, other in prompt_histories.items():
|
||||
if other_key == prompt_key:
|
||||
# ====================== WORD GROUPING FROM PROMPT ======================
|
||||
word_groups = [] # (start, end, full_word, word_token_list)
|
||||
seen = set()
|
||||
if tokenizer and positive_prompt:
|
||||
raw_words = [w.strip() for w in positive_prompt.split() if w.strip()]
|
||||
for word in raw_words:
|
||||
clean_word = word
|
||||
lower = clean_word.lower()
|
||||
if lower in seen or len(clean_word) < 3 or not self._is_valuable_token(clean_word):
|
||||
continue
|
||||
other_tokens = set()
|
||||
for entry in other.get("history", []):
|
||||
if "prompt" in entry and entry["prompt"]:
|
||||
try:
|
||||
other_tokens.update(tokenizer.encode(entry["prompt"], add_special_tokens=True)[:60])
|
||||
except:
|
||||
pass
|
||||
if not other_tokens:
|
||||
continue
|
||||
intersection = len(current_set & other_tokens)
|
||||
union = len(current_set | other_tokens)
|
||||
sim_weight = intersection / union if union > 0 else 0.0
|
||||
if sim_weight < 0.3:
|
||||
continue
|
||||
for entry in other.get("history", []):
|
||||
if "prompt" not in entry or not entry["prompt"]:
|
||||
continue
|
||||
try:
|
||||
old_tokens = tokenizer.encode(entry["prompt"], add_special_tokens=True)[:60]
|
||||
matching = current_set & set(old_tokens)
|
||||
r = entry.get("rating", 5)
|
||||
adj = 1.25 if r >= 8 else 0.75 if r <= 3 else None
|
||||
if adj is not None:
|
||||
for tid in matching:
|
||||
tid_str = str(tid)
|
||||
if tid_str in token_importance:
|
||||
token_importance[tid_str] = max(0.3, min(2.5, token_importance[tid_str] * (1.0 + (adj - 1.0) * sim_weight)))
|
||||
except:
|
||||
pass
|
||||
seen.add(lower)
|
||||
|
||||
word_token_list = tokenizer.encode(clean_word, add_special_tokens=False)
|
||||
|
||||
found = False
|
||||
for start in range(len(full_token_ids) - len(word_token_list) + 1):
|
||||
if full_token_ids[start:start + len(word_token_list)] == word_token_list:
|
||||
word_groups.append((start, start + len(word_token_list), clean_word, word_token_list))
|
||||
found = True
|
||||
break
|
||||
if not found and word_token_list:
|
||||
word_groups.append((0, 1, clean_word, [word_token_list[0]]))
|
||||
|
||||
# ====================== SCHEDULER SETUP ======================
|
||||
global_adaptive["scheduler_mode"] = scheduler_mode
|
||||
@@ -388,86 +347,34 @@ class FunPackGemmaEmbeddingRefiner:
|
||||
cur_mean = cur_positive.mean(dim=1)
|
||||
similarity = F.cosine_similarity(ref_mean, cur_mean, dim=-1).mean().item()
|
||||
reward = (rating - 5.5) / 4.5
|
||||
last_rating = active.get("last_rating", 5)
|
||||
last_reward = (last_rating - 5.5) / 4.5
|
||||
|
||||
lr_scale, confidence, exploration_mult, token_lr_mult = self._get_scheduler_factors(
|
||||
word_importance = global_adaptive.setdefault("word_importance", {})
|
||||
|
||||
lr_scale, confidence, exploration_mult, word_lr_mult = self._get_scheduler_factors(
|
||||
scheduler_mode, rating, reward, similarity, iter_num, total_iters,
|
||||
global_adaptive["token_importance"], word_groups, full_token_ids, global_adaptive, device
|
||||
word_importance, word_groups, global_adaptive, device
|
||||
)
|
||||
|
||||
expl = global_adaptive["exploration_base"] * exploration_mult
|
||||
|
||||
# ====================== TOKEN LIBRARY UPDATE ======================
|
||||
token_library = global_adaptive.setdefault("token_library", {})
|
||||
token_importance = global_adaptive["token_importance"]
|
||||
infusion_tokens = global_adaptive["infusion_tokens"]
|
||||
# ====================== WORD IMPORTANCE UPDATE ======================
|
||||
for start, end, full_word, word_token_list in word_groups:
|
||||
if not self._is_valuable_token(full_word):
|
||||
continue
|
||||
wkey = full_word.lower()
|
||||
if wkey not in word_importance:
|
||||
word_importance[wkey] = 1.0
|
||||
|
||||
if tokenizer and full_token_ids:
|
||||
valuable_token_ids = [tid for tid in full_token_ids if self._is_valuable_token(tokenizer.decode([tid]))]
|
||||
norm_factor = len(valuable_token_ids) if valuable_token_ids else 1
|
||||
for tid in full_token_ids:
|
||||
tid_str = str(tid)
|
||||
if tid_str not in token_library:
|
||||
token_library[tid_str] = {
|
||||
"fqc": 0,
|
||||
"unique_prompts": 0,
|
||||
"high_rated_count": 0,
|
||||
"low_rated_count": 0,
|
||||
"good_pairs": {},
|
||||
"contrib_score": 0.0,
|
||||
"redundancy_score": 0.0,
|
||||
"protected_until": 0
|
||||
}
|
||||
lib_entry = token_library[tid_str]
|
||||
lib_entry["fqc"] += 1
|
||||
if is_new_prompt:
|
||||
lib_entry["unique_prompts"] += 1
|
||||
is_new = True
|
||||
base_lr = 0.22 / (1 + 0.07 * (len(history) ** 0.5))
|
||||
group_delta = reward * base_lr * lr_scale * confidence * (1.8 if is_new else 1.0)
|
||||
|
||||
if rating >= 8:
|
||||
for i, tid1 in enumerate(full_token_ids):
|
||||
tid1_str = str(tid1)
|
||||
for tid2 in full_token_ids[i+1:]:
|
||||
tid2_str = str(tid2)
|
||||
if tid1_str in token_library and tid2_str in token_library:
|
||||
if tid2_str not in token_library[tid1_str]["good_pairs"]:
|
||||
token_library[tid1_str]["good_pairs"][tid2_str] = 0
|
||||
if tid1_str not in token_library[tid2_str]["good_pairs"]:
|
||||
token_library[tid2_str]["good_pairs"][tid1_str] = 0
|
||||
token_library[tid1_str]["good_pairs"][tid2_str] += 1
|
||||
token_library[tid2_str]["good_pairs"][tid1_str] += 1
|
||||
if tid1_str in token_library:
|
||||
token_library[tid1_str]["high_rated_count"] += 1
|
||||
elif rating <= 4:
|
||||
for tid in full_token_ids:
|
||||
tid_str = str(tid)
|
||||
if tid_str in token_library:
|
||||
token_library[tid_str]["low_rated_count"] += 1
|
||||
if wkey in word_lr_mult:
|
||||
group_delta *= word_lr_mult[wkey]
|
||||
|
||||
# Safe contextual updates
|
||||
for start, end, full_word, word_token_list in word_groups:
|
||||
if not self._is_valuable_token(full_word):
|
||||
continue
|
||||
tid_str = str(word_token_list[0])
|
||||
entry = token_library[tid_str]
|
||||
entry["contrib_score"] = 0.85 * entry.get("contrib_score", 0.0) + 0.15 * 0.6
|
||||
if len(word_groups) > 4:
|
||||
entry["redundancy_score"] = min(1.0, entry.get("redundancy_score", 0.0) + 0.07)
|
||||
if rating >= 8 and entry.get("contrib_score", 0) > 0.3:
|
||||
entry["protected_until"] = max(entry.get("protected_until", 0), iter_num + 8)
|
||||
word_importance[wkey] = max(0.35, min(2.8, word_importance[wkey] + group_delta))
|
||||
|
||||
for tid_str, entry in token_library.items():
|
||||
total_rated = entry["high_rated_count"] + entry["low_rated_count"]
|
||||
if total_rated == 0:
|
||||
continue
|
||||
appears_in_many = entry["unique_prompts"] >= 5
|
||||
balanced = entry["high_rated_count"] > 0 and entry["low_rated_count"] > 0
|
||||
mostly_high = entry["high_rated_count"] > entry["low_rated_count"] * 2
|
||||
entry["is_general"] = appears_in_many and balanced
|
||||
entry["is_specific"] = not entry["is_general"] and mostly_high and entry["unique_prompts"] <= 8
|
||||
entry["infusable"] = (entry["is_general"] or entry["is_specific"]) and entry["high_rated_count"] >= 2
|
||||
|
||||
# ====================== CORE REFINEMENT ======================
|
||||
# ====================== CORE REFINEMENT (POSITIVE ONLY) ======================
|
||||
reference = cur_positive.clone()
|
||||
momentum = global_adaptive.get("momentum")
|
||||
if momentum is None or not isinstance(momentum, dict):
|
||||
@@ -493,91 +400,11 @@ class FunPackGemmaEmbeddingRefiner:
|
||||
sim_threshold = global_adaptive["dynamic_sim_threshold"]
|
||||
is_close = (not is_new_prompt) and (similarity >= sim_threshold)
|
||||
|
||||
new_token_mask = None
|
||||
infusion_active = False
|
||||
infusion_weights = None
|
||||
|
||||
if tokenizer and positive_prompt:
|
||||
try:
|
||||
embed_seq_len = cur_positive.shape[1] if cur_positive.dim() == 3 else cur_positive.shape[0] if cur_positive.dim() == 2 else 1
|
||||
prev_token_ids = tokenizer.encode(active.get("last_positive_prompt", ""), add_special_tokens=True) if "last_positive_prompt" in active else []
|
||||
new_token_mask = torch.ones(embed_seq_len, device=device)
|
||||
compare_len = min(len(full_token_ids), len(prev_token_ids), embed_seq_len)
|
||||
for i in range(compare_len):
|
||||
if full_token_ids[i] != prev_token_ids[i]:
|
||||
new_token_mask[i] = 0.0
|
||||
|
||||
# Word-level importance update with scheduler
|
||||
for start, end, full_word, word_token_list in word_groups:
|
||||
if not self._is_valuable_token(full_word):
|
||||
continue
|
||||
group_key = str(word_token_list[0])
|
||||
if group_key not in token_importance:
|
||||
token_importance[group_key] = 1.0
|
||||
|
||||
is_new = any(tid not in prev_token_ids for tid in word_token_list)
|
||||
update_strength = 1.8 if is_new else 1.0
|
||||
|
||||
# Missing token protection
|
||||
missing_penalty = 1.0
|
||||
if group_key in token_importance and group_key not in [str(t) for t in full_token_ids]:
|
||||
if token_library.get(group_key, {}).get("contrib_score", 0) > 0.35:
|
||||
missing_penalty = 0.78
|
||||
else:
|
||||
missing_penalty = 0.55
|
||||
|
||||
fqc = token_library.get(group_key, {}).get("fqc", 1)
|
||||
base_lr = 0.22 / (1 + 0.07 * (fqc ** 0.5))
|
||||
|
||||
group_delta = reward * base_lr * lr_scale * confidence * update_strength * missing_penalty
|
||||
|
||||
for tid in word_token_list:
|
||||
tid_str = str(tid)
|
||||
if tid_str in token_lr_mult:
|
||||
group_delta *= token_lr_mult[tid_str]
|
||||
token_importance[tid_str] = max(0.35, min(2.8, token_importance[tid_str] + group_delta))
|
||||
|
||||
# Infusion logic (kept from original)
|
||||
if random.random() < 0.18 and token_library:
|
||||
infusion_weights = torch.ones(embed_seq_len, device=device)
|
||||
infusion_any = False
|
||||
for i, tid in enumerate(full_token_ids[:embed_seq_len]):
|
||||
tid_str = str(tid)
|
||||
if tid_str not in token_library or not token_library[tid_str].get("infusable", False):
|
||||
continue
|
||||
token_text = tokenizer.decode([int(tid)], skip_special_tokens=True).strip()
|
||||
if not self._is_valuable_token(token_text):
|
||||
continue
|
||||
entry = token_library[tid_str]
|
||||
if not entry.get("good_pairs"):
|
||||
continue
|
||||
pair_scores = [entry["good_pairs"].get(str(ct), 0) for ct in full_token_ids[:embed_seq_len]]
|
||||
pair_scores = [s for s in pair_scores if s > 0]
|
||||
if pair_scores:
|
||||
avg_pair = sum(pair_scores) / len(pair_scores)
|
||||
if avg_pair > 1.2 or entry.get("is_general", False):
|
||||
if tid_str not in infusion_tokens:
|
||||
infusion_tokens[tid_str] = 0.0
|
||||
infusion_tokens[tid_str] = min(3.0, infusion_tokens[tid_str] + 0.6)
|
||||
infusion_weights[i] += infusion_tokens[tid_str] * 0.15
|
||||
infusion_any = True
|
||||
if infusion_any:
|
||||
infusion_active = True
|
||||
except Exception as e:
|
||||
print(f"[FunPackGemmaEmbeddingRefiner] Token comparison failed: {e}")
|
||||
new_token_mask = None
|
||||
infusion_weights = None
|
||||
|
||||
# Delta calculation
|
||||
# Simple new_delta
|
||||
if is_close and history:
|
||||
last_entry = history[-1]
|
||||
mod_data = last_entry.get("modified_embeds")
|
||||
if mod_data is None:
|
||||
prev_modified = torch.zeros_like(old_reference)
|
||||
else:
|
||||
prev_modified = serializable_to_tensor(mod_data).to(device)
|
||||
if list(prev_modified.shape) != list(old_reference.shape):
|
||||
prev_modified = torch.zeros_like(old_reference)
|
||||
prev_modified = serializable_to_tensor(mod_data).to(device) if mod_data else torch.zeros_like(old_reference)
|
||||
prev_delta = prev_modified - old_reference
|
||||
multiplier = max(0.05, 1.0 + reward * 1.45)
|
||||
if rating == last_rating and rating >= 8:
|
||||
@@ -588,96 +415,28 @@ class FunPackGemmaEmbeddingRefiner:
|
||||
good_deltas = []
|
||||
for entry in history:
|
||||
if entry.get("rating", 0) >= 7:
|
||||
mod_data = entry.get("modified_embeds")
|
||||
if mod_data is None:
|
||||
continue
|
||||
mod = serializable_to_tensor(mod_data).to(device)
|
||||
mod = serializable_to_tensor(entry.get("modified_embeds")).to(device)
|
||||
if list(mod.shape) == list(old_reference.shape):
|
||||
good_deltas.append(mod - old_reference)
|
||||
if good_deltas:
|
||||
new_delta = torch.stack(good_deltas).mean(dim=0) * (0.7 + reward * 0.4)
|
||||
else:
|
||||
if is_new_prompt and isinstance(momentum, torch.Tensor) and list(momentum.shape) == list(reference.shape):
|
||||
seed_strength = min(0.3, abs(avg_reward_ema) * 0.4)
|
||||
new_delta = momentum * seed_strength + torch.randn_like(reference) * expl * 0.3
|
||||
else:
|
||||
new_delta = torch.randn_like(reference) * expl * 0.45
|
||||
new_delta = torch.randn_like(reference) * expl * 0.45
|
||||
|
||||
# Apply importance, new token boost, infusion
|
||||
if token_importance and cur_positive.dim() > 1:
|
||||
# Apply WORD-LEVEL importance
|
||||
if word_importance and cur_positive.dim() > 1:
|
||||
seq_len = cur_positive.shape[1] if cur_positive.dim() == 3 else cur_positive.shape[0]
|
||||
importance_tensor = torch.ones(seq_len, device=device)
|
||||
for i, tid in enumerate(full_token_ids[:seq_len]):
|
||||
tid_str = str(tid)
|
||||
if tid_str in token_importance:
|
||||
importance_tensor[i] = token_importance[tid_str]
|
||||
for start, end, full_word, _ in word_groups:
|
||||
imp = word_importance.get(full_word.lower(), 1.0)
|
||||
importance_tensor[start:end] = imp
|
||||
new_delta = new_delta * importance_tensor.unsqueeze(-1)
|
||||
|
||||
if new_token_mask is not None:
|
||||
boost_factor = 1.0 + 1.8 * (1.0 - new_token_mask).unsqueeze(-1)
|
||||
new_delta = new_delta * boost_factor
|
||||
|
||||
if infusion_active and infusion_weights is not None:
|
||||
new_delta = new_delta * infusion_weights.unsqueeze(-1)
|
||||
|
||||
new_positive = reference + new_delta
|
||||
new_positive = torch.clamp(new_positive, min=-60.0, max=60.0)
|
||||
norm_factor = reference.norm(dim=-1, keepdim=True) + 1e-8
|
||||
new_positive = new_positive / (new_positive.norm(dim=-1, keepdim=True) + 1e-8) * norm_factor
|
||||
|
||||
# ====================== NEGATIVE REFINEMENT ======================
|
||||
new_negative = None
|
||||
negative_meta = {"pooled_output": None}
|
||||
if negative_conditioning is not None:
|
||||
neg_item = negative_conditioning[0]
|
||||
if isinstance(neg_item, (list, tuple)) and len(neg_item) >= 2:
|
||||
raw_negative = neg_item[0]
|
||||
negative_meta = neg_item[1] if isinstance(neg_item[1], dict) else {"pooled_output": None}
|
||||
else:
|
||||
raw_negative = neg_item if isinstance(neg_item, torch.Tensor) else None
|
||||
if isinstance(raw_negative, torch.Tensor):
|
||||
raw_negative = raw_negative.to(device)
|
||||
neg_strength = 0.15 if rating <= 4 else 0.05
|
||||
neg_delta = torch.randn_like(raw_negative) * neg_strength * (5.5 - rating) / 4.5
|
||||
new_negative = raw_negative + neg_delta
|
||||
new_negative = torch.clamp(new_negative, min=-60.0, max=60.0)
|
||||
norm_factor_neg = raw_negative.norm(dim=-1, keepdim=True) + 1e-8
|
||||
new_negative = new_negative / (new_negative.norm(dim=-1, keepdim=True) + 1e-8) * norm_factor_neg
|
||||
|
||||
# ====================== LATENT REFINEMENT ======================
|
||||
modified_latent = latent
|
||||
latent_info = " | Latent unchanged"
|
||||
if latent is not None and "samples" in latent:
|
||||
samples = latent["samples"].to(device) if latent["samples"].device != device else latent["samples"]
|
||||
current_shape = list(samples.shape)
|
||||
good_latents = []
|
||||
bad_latents = []
|
||||
for entry in history:
|
||||
if "latent_samples" in entry:
|
||||
mod_data = entry["latent_samples"]
|
||||
if mod_data is not None:
|
||||
mod = serializable_to_tensor(mod_data).to(device)
|
||||
if list(mod.shape) == current_shape:
|
||||
if entry.get("rating", 5) >= 8:
|
||||
good_latents.append(mod)
|
||||
elif entry.get("rating", 5) <= 3:
|
||||
bad_latents.append(mod)
|
||||
if good_latents or bad_latents:
|
||||
latent_strength = max(0.0, avg_reward_ema * 0.35)
|
||||
latent_delta = torch.zeros_like(samples)
|
||||
if good_latents:
|
||||
mean_good = torch.stack(good_latents).mean(dim=0)
|
||||
latent_delta += (mean_good - samples) * latent_strength
|
||||
if bad_latents:
|
||||
mean_bad = torch.stack(bad_latents).mean(dim=0)
|
||||
latent_delta -= (mean_bad - samples) * (latent_strength * 0.65)
|
||||
mask = torch.abs(samples) < 0.05
|
||||
latent_delta = latent_delta * (~mask).float()
|
||||
if torch.abs(latent_delta).max() > 1e-4:
|
||||
new_samples = torch.clamp(samples + latent_delta, min=-12.0, max=12.0)
|
||||
modified_latent = {"samples": new_samples.cpu(), "noise_mask": latent.get("noise_mask")}
|
||||
latent_info = f" | Latent ±{len(good_latents)}g/{len(bad_latents)}b"
|
||||
|
||||
# Update momentum
|
||||
momentum = 0.75 * momentum + 0.25 * (new_delta * reward)
|
||||
global_adaptive["momentum"] = tensor_to_serializable(momentum)
|
||||
@@ -696,8 +455,6 @@ class FunPackGemmaEmbeddingRefiner:
|
||||
"similarity": round(similarity, 4),
|
||||
"prompt": positive_prompt[:180]
|
||||
}
|
||||
if latent is not None and "samples" in latent:
|
||||
history_entry["latent_samples"] = tensor_to_serializable(latent["samples"])
|
||||
history.append(history_entry)
|
||||
|
||||
if not unlimited_history and len(history) > 200:
|
||||
@@ -715,107 +472,63 @@ class FunPackGemmaEmbeddingRefiner:
|
||||
data["global_adaptive"] = global_adaptive
|
||||
|
||||
# ====================== INTELLIGENT FEEDBACK QUESTION ======================
|
||||
if feedback_enabled and data.get("pending_feedback") is None and tokenizer and word_groups:
|
||||
if feedback_enabled and data.get("pending_feedback") is None and word_groups:
|
||||
feedbacked = set(global_adaptive.get("feedbacked_words", []))
|
||||
last_word = global_adaptive.get("last_feedback_word")
|
||||
new_unasked = []
|
||||
for start, end, full_word, word_token_list in word_groups:
|
||||
if not self._is_valuable_token(full_word):
|
||||
candidates = []
|
||||
for _, _, full_word, word_token_list in word_groups:
|
||||
if not self._is_valuable_token(full_word) or full_word.lower() in feedbacked:
|
||||
continue
|
||||
if full_word.lower() in feedbacked or full_word == last_word:
|
||||
continue
|
||||
group_key = str(word_token_list[0])
|
||||
imp = token_importance.get(group_key, 1.0)
|
||||
freq = token_library.get(group_key, {}).get("fqc", 1)
|
||||
score = freq * (1.0 - abs(imp - 1.0))
|
||||
new_unasked.append((full_word, group_key, score, word_token_list[0]))
|
||||
if new_unasked:
|
||||
chosen = max(new_unasked, key=lambda x: x[2])
|
||||
else:
|
||||
candidates = []
|
||||
for start, end, full_word, word_token_list in word_groups:
|
||||
if not self._is_valuable_token(full_word) or full_word == last_word:
|
||||
continue
|
||||
group_key = str(word_token_list[0])
|
||||
imp = token_importance.get(group_key, 1.0)
|
||||
freq = token_library.get(group_key, {}).get("fqc", 1)
|
||||
score = freq * (1.0 - abs(imp - 1.0))
|
||||
rating_bonus = 1.5 if rating >= 8 or rating <= 3 else 1.0
|
||||
candidates.append((full_word, group_key, score * rating_bonus, word_token_list[0]))
|
||||
if candidates:
|
||||
chosen = max(candidates, key=lambda x: x[2])
|
||||
else:
|
||||
chosen = None
|
||||
if chosen:
|
||||
full_word, group_key, _, best_tid = chosen
|
||||
wkey = full_word.lower()
|
||||
imp = word_importance.get(wkey, 1.0)
|
||||
score = 1.0 - abs(imp - 1.0)
|
||||
candidates.append((full_word, wkey, score, word_token_list))
|
||||
|
||||
if candidates:
|
||||
chosen = max(candidates, key=lambda x: x[2])
|
||||
full_word, wkey, _, word_token_list = chosen
|
||||
|
||||
data["pending_feedback"] = {
|
||||
"token_id": int(best_tid),
|
||||
"token_text": full_word,
|
||||
"full_word": full_word,
|
||||
"word_token_list": word_token_list,
|
||||
"prompt_key": prompt_key,
|
||||
"iteration": iter_num
|
||||
}
|
||||
global_adaptive["last_feedback_tid"] = int(best_tid)
|
||||
global_adaptive["last_feedback_word"] = full_word
|
||||
if full_word.lower() not in feedbacked:
|
||||
global_adaptive.setdefault("feedbacked_words", []).append(full_word)
|
||||
global_adaptive.setdefault("feedbacked_words", []).append(full_word)
|
||||
|
||||
feedback_question_output = f"Word: '{full_word}' – How well represented? (1=absent, 5=perfect, 6=overrepresented)"
|
||||
else:
|
||||
feedback_question_output = "Feedback enabled but no suitable word found in current prompt."
|
||||
|
||||
# ====================== TRAINING INFO ======================
|
||||
current_top = self._get_top_tokens(token_importance, tokenizer, 12)
|
||||
top_infusion = self._get_top_tokens(infusion_tokens, tokenizer, 8)
|
||||
current_top = self._get_top_tokens(word_importance, tokenizer, 12)
|
||||
|
||||
top_contrib = sorted(
|
||||
[(k, v.get("contrib_score", 0.0)) for k, v in token_library.items() if v.get("contrib_score", 0.0) > 0.1],
|
||||
key=lambda x: x[1], reverse=True
|
||||
)[:10]
|
||||
top_contrib_str = ", ".join([
|
||||
f"{tokenizer.decode([int(k)], skip_special_tokens=True).strip()}({v:.2f})"
|
||||
for k, v in top_contrib
|
||||
]) or "None"
|
||||
normalized_reward = max(0.0, min(1.0, (avg_reward_ema + 1.0) / 2.0))
|
||||
stability_factor = max(0.0, 1.0 - similarity)
|
||||
raw_loss = (1.0 - normalized_reward) * 0.7 + stability_factor * 0.3
|
||||
learning_loss = max(0.02, raw_loss * (1.0 - min(0.95, good_ratio * 0.8)))
|
||||
|
||||
training_info = (
|
||||
f"Scheduler: {scheduler_mode.upper()} | Global Step: {global_adaptive['current_step']} | "
|
||||
f"Scheduler: {scheduler_mode.upper()} | Step: {global_adaptive['current_step']} | "
|
||||
f"EMA Reward: {avg_reward_ema:+.3f} | Confidence: {confidence:.2f} | LR Scale: {lr_scale:.3f}\n"
|
||||
f"Exploration: {expl:.3f} | Similarity: {similarity:.4f} | Raw Reward: {reward:+.3f} | Good Ratio: {good_ratio:.1%}\n"
|
||||
f"Top Importance Tokens: {current_top}\n"
|
||||
f"Top Contributors (by delta impact): {top_contrib_str}\n"
|
||||
f"Protected: {sum(1 for v in token_library.values() if v.get('protected_until', 0) > iter_num)} | "
|
||||
f"Library: {len(token_library)} tokens | Infusion active: {infusion_active}"
|
||||
f"Exploration: {expl:.3f} | Similarity: {similarity:.4f} | Good Ratio: {good_ratio:.1%}\n"
|
||||
f"**Learning Loss: {learning_loss:.4f}** (lower is better)\n"
|
||||
f"Top Word Importance: {current_top}\n"
|
||||
f"Word Library Size: {len(word_importance)}"
|
||||
)
|
||||
|
||||
if scheduler_mode == "accurate":
|
||||
training_info += "\n[Accurate] Conservative Prodigy + Cosine decay - slow & stable"
|
||||
training_info += "\n[Accurate] Conservative • Prodigy + Cosine"
|
||||
elif scheduler_mode == "aggressive":
|
||||
training_info += "\n[Aggressive] Fast early boosts - quick style locking"
|
||||
training_info += "\n[Aggressive] Fast style locking"
|
||||
|
||||
# ====================== STATUS ======================
|
||||
iter_num = len(history)
|
||||
total_iterations = sum(len(p.get("history", [])) for p in prompt_histories.values())
|
||||
unique_prompts = len(prompt_histories)
|
||||
infusion_count = len([v for v in infusion_tokens.values() if v > 0.5])
|
||||
trend = "↑" if reward > last_reward else "↓" if reward < last_reward else "→"
|
||||
health = "🚀 Strong convergence" if avg_reward_ema > 0.6 else "✅ Learning well" if avg_reward_ema > 0.3 else "⚠️ Still exploring" if avg_reward_ema > -0.2 else "🔄 Heavy correction"
|
||||
|
||||
status = (f"Iter {iter_num} | Total iters {total_iterations} | "
|
||||
f"History {len(history)} | Unique prompts {unique_prompts} | "
|
||||
f"Infusion tokens {infusion_count}\n"
|
||||
status = (f"Iter {iter_num} | Total iters {total_iters} | "
|
||||
f"History {len(history)} | Unique prompts {len(prompt_histories)}\n"
|
||||
f"Rating {rating}/10 {trend} | Sim {similarity:.3f} | Reward {reward:+.2f} | "
|
||||
f"EMA {avg_reward_ema:+.2f} | Good ratio {good_ratio:.0%} | Expl {expl:.3f} | {health}\n"
|
||||
f"Current focus: {current_top}\n"
|
||||
f"Top infusion: {top_infusion}\n"
|
||||
f"Library: {len(token_library)} tokens")
|
||||
|
||||
if is_new_prompt:
|
||||
status += "\n✓ New prompt pair → token overlap transfer applied"
|
||||
if new_token_mask is not None and (1.0 - new_token_mask).sum() > 0:
|
||||
status += "\nNew tokens boosted"
|
||||
if infusion_active:
|
||||
status += "\n🔀 Token infusion active (delta scaled)"
|
||||
if rating <= 4:
|
||||
status += "\nLow rating: importance down-weighted + negative strengthened"
|
||||
status += latent_info
|
||||
f"Current focus: {current_top}")
|
||||
|
||||
# ====================== FINAL SAVE ======================
|
||||
with open(json_file, "w", encoding="utf-8") as f:
|
||||
@@ -823,9 +536,8 @@ class FunPackGemmaEmbeddingRefiner:
|
||||
|
||||
# ====================== RETURN ======================
|
||||
modified_positive = [(new_positive, positive_meta)]
|
||||
modified_negative = [(new_negative, negative_meta)] if new_negative is not None else negative_conditioning
|
||||
|
||||
return (modified_positive, modified_negative, modified_latent, status, feedback_question_output, training_info)
|
||||
return (modified_positive, status, feedback_question_output, training_info)
|
||||
|
||||
# Constants from StoryMem
|
||||
IMAGE_FACTOR = 28
|
||||
|
||||
Reference in New Issue
Block a user