39 lines
1.2 KiB
Python
39 lines
1.2 KiB
Python
import torch
|
|
from torch.nn import Embedding
|
|
|
|
from custom_nodes.KepPromptLang.lib.action.base import (
|
|
Action,
|
|
SingleArgAction,
|
|
)
|
|
from custom_nodes.KepPromptLang.lib.parser.registration import register_action
|
|
|
|
|
|
class NormAction(SingleArgAction):
|
|
grammar = 'norm(" arg+ ")"'
|
|
name = "norm"
|
|
chars = None
|
|
|
|
def token_length(self) -> int:
|
|
"""
|
|
Norm normalizes the embeddings of the base segment, so the length is the length of the base segment
|
|
:return:
|
|
"""
|
|
total_length = 0
|
|
for seg_or_action in self.arg:
|
|
total_length += seg_or_action.token_length()
|
|
|
|
return total_length
|
|
|
|
def get_result(self, embedding_module: Embedding) -> torch.Tensor:
|
|
all_embeddings = []
|
|
for seg_or_action in self.arg:
|
|
if isinstance(seg_or_action, Action):
|
|
target_embeddings = seg_or_action.get_result(embedding_module)
|
|
else:
|
|
target_embeddings = seg_or_action.get_embeddings(embedding_module)
|
|
all_embeddings.append(target_embeddings)
|
|
|
|
target_embeddings = torch.cat(all_embeddings, dim=1)
|
|
return torch.div(target_embeddings, torch.norm(target_embeddings, dim=-1, keepdim=True))
|
|
|