diff --git a/lib/__init__.py b/lib/__init__.py index 6ff6fe8..2e6ad4f 100644 --- a/lib/__init__.py +++ b/lib/__init__.py @@ -1,3 +1,4 @@ +from custom_nodes.KepPromptLang.lib.actions.abs_max import AbsMaxAction from custom_nodes.KepPromptLang.lib.actions.avg import AverageAction from custom_nodes.KepPromptLang.lib.actions.diff import DiffAction from custom_nodes.KepPromptLang.lib.actions.max import MaxAction @@ -22,3 +23,4 @@ register_action(AverageAction) register_action(ScaleDims) register_action(SetDims) register_action(MaxAction) +register_action(AbsMaxAction) diff --git a/lib/actions/abs_max.py b/lib/actions/abs_max.py new file mode 100644 index 0000000..4c5f379 --- /dev/null +++ b/lib/actions/abs_max.py @@ -0,0 +1,104 @@ +from typing import List, Union + +import torch +from torch.nn import Embedding + +from custom_nodes.KepPromptLang.lib.action.base import Action, MultiArgAction +from custom_nodes.KepPromptLang.lib.actions.action_utils import get_embedding +from custom_nodes.KepPromptLang.lib.actions.utils import is_broadcastable +from custom_nodes.KepPromptLang.lib.parser.prompt_segment import PromptSegment + + +class AbsMaxAction(MultiArgAction): + grammar = 'absMax(" arg ("|" arg)+ ")"' + name = "absMax" + chars = ["+", "+"] + + def __init__(self, args: List[List[Union[PromptSegment, Action]]]) -> None: + super().__init__(args) + + self.base_arg = args[0] + self.additional_args = args[1:] + + def token_length(self) -> int: + # AbsMax modifies the base segment, so the length is the length of the base segment + return sum(seg_or_action.token_length() for seg_or_action in self.base_arg) + + def get_result(self, embedding_module: Embedding) -> torch.Tensor: + # Calculate the embeddings for the base segment + all_base_embeddings = [ + get_embedding(seg_or_action, embedding_module) + for seg_or_action in self.base_arg + ] + + result = torch.cat(all_base_embeddings, dim=1) + + for arg in self.additional_args: + all_arg_embeddings = [ + get_embedding(seg_or_action, embedding_module) for seg_or_action in arg + ] + + arg_embedding = torch.cat(all_arg_embeddings, dim=1) + + if is_broadcastable(result, arg_embedding): + positive_mask1 = result > 0 + positive_mask2 = arg_embedding > 0 + negative_mask1 = result < 0 + negative_mask2 = arg_embedding < 0 + zero_mask1 = result == 0 + zero_mask2 = arg_embedding == 0 + + # For mixed signs, choose the one with the largest magnitude + mixed_mask_pos_neg = positive_mask1 & negative_mask2 + mixed_mask_neg_pos = negative_mask1 & positive_mask2 + mixed_mask = mixed_mask_pos_neg | mixed_mask_neg_pos + mixed_selection = torch.where( + result.abs() > arg_embedding.abs(), result, arg_embedding + ) + + # Apply max for positive dimensions, min for negative dimensions, and handle zeros and mixed signs + result = torch.where( + positive_mask1 & positive_mask2, + torch.max(result, arg_embedding), + torch.where( + negative_mask1 & negative_mask2, + torch.min(result, arg_embedding), + torch.where( + positive_mask1 & zero_mask2, + result, + torch.where( + zero_mask1 & positive_mask2, + arg_embedding, + torch.where(mixed_mask, mixed_selection, result), + ), + ), + ), + ) + else: + print( + "WARNING: shape mismatch when trying to apply absMax, arg will be averaged" + ) + result = torch.max(result, torch.mean(arg_embedding, dim=1, keepdim=True)) + return result + + # def __repr__(self): + # return f"sum(\n\tbase_segment={self.base_segment},\n\targs={self.args}\n)" + def __repr__(self) -> str: + return f"sum({', '.join(map(str, self.additional_args))})" + + def depth_repr(self, depth=1): + out = "NudgeAction(\n" + if isinstance(self.base_arg, Action): + base_segment_repr = self.base_arg.depth_repr(depth + 1) + out += "\t" * depth + f"base_segment={base_segment_repr}\n" + else: + out += "\t" * depth + f"base_segment={self.base_arg.depth_repr()},\n" + + if isinstance(self.additional_args, Action): + target_repr = self.additional_args.depth_repr(depth + 1) + out += "\t" * depth + f"target={target_repr},\n" + else: + out += "\t" * depth + f"target={self.additional_args.depth_repr()},\n" + out += "\t" * depth + f"weight={self.weight},\n" + out += "\t" * (depth - 1) + ")" + return out