psaiops 0.2.0__py3-none-any.whl → 0.2.2__py3-none-any.whl

psaiops/compose/contrast/app.py

@@ -1,6 +1,7 @@
 import functools
 
 import gradio
+import pandas
 import torch
 import torch.cuda
 
@@ -39,7 +40,7 @@ def create_model_block() -> dict:
 # SAMPLING #####################################################################
 
 def create_sampling_block() -> dict:
-    __tokens = gradio.Slider(label='Tokens', value=16, minimum=1, maximum=128, step=1, scale=1, interactive=True)
+    __tokens = gradio.Slider(label='Tokens', value=32, minimum=1, maximum=128, step=1, scale=1, interactive=True)
     __topk = gradio.Slider(label='Top K', value=4, minimum=1, maximum=8, step=1, scale=1, interactive=True)
     __topp = gradio.Slider(label='Top P', value=0.9, minimum=0.0, maximum=1.0, step=0.1, scale=1, interactive=True)
     return {
@@ -129,9 +130,22 @@ def create_layout(intro: str=INTRO) -> dict:
 def update_layer_range(value: float, model: str) -> dict:
     return gradio.update(maximum=35, value=min(35, int(value))) if '120b' in model else gradio.update(maximum=23, value=min(23, int(value)))
 
-def update_table_data(positive: str, negative: str, prompt: str, output: str, tokenizer: object) -> list:
-    __outputs = tokenizer([positive, negative, prompt, output], padding=True)
-    return [tokenizer.convert_ids_to_tokens(__s) for __s in __outputs['input_ids']]
+def update_table_data(positive: str, negative: str, prompt: str, output: str, tokenizer: object) -> pandas.DataFrame:
+    # array of token IDs
+    __outputs = tokenizer([positive, negative, prompt, output], return_tensors='pt', padding=True)
+    # array of token strings
+    __tokens = [tokenizer.convert_ids_to_tokens(__s) for __s in __outputs['input_ids']]
+    # shift the special characters
+    __tokens = [[__t.replace(chr(0x0120), ' ').replace(chr(0x010a), '\\n') for __t in __s] for __s in __tokens]
+    # mask the tokens that differ between positive and negative prompts
+    __masks = psaiops.compose.contrast.lib.compute_sequence_mask(tokens=__outputs['input_ids'])
+    # convert into a data frame
+    __data = pandas.DataFrame(__tokens)
+    # color the background in red for the positions marked by the mask
+    return __data.style.apply(update_table_style, masks=pandas.DataFrame(__masks), axis=None)
+
+def update_table_style(data: pandas.DataFrame, masks: pandas.DataFrame) -> pandas.DataFrame:
+    return pandas.DataFrame(masks.replace({True: 'background-color: rgb(255, 0, 0, 64%)', False: 'background-color: rgb(0, 0, 0, 0%)',}))
 
 # APP ##########################################################################
 

psaiops/compose/contrast/lib.py

@@ -57,33 +57,37 @@ def capture_hidden_activation(
 
 def compute_sequence_mask(
     tokens: torch.Tensor, # (B, S)
-    masks: torch.Tensor, # (B, S)
+    # masks: torch.Tensor, # (B, S)
 ) -> torch.Tensor:
-    # keep the intersection of the masks
-    __masks = masks.prod(dim=0, keepdim=False)
-    # and only the positions that differ between positive and negative prompts
-    return __masks * (tokens[0] != tokens[1])
+    __shape = mlable.shapes.divide(tokens.shape, axis=0, factor=2, insert=True)
+    # group the samples two by two
+    __data = tokens.reshape(__shape)
+    # compare each sample with its neighbor
+    __masks = __data[:, :1]  != __data[:, 1:]
+    # apply the same mask to both samples
+    return __masks.expand(__shape).reshape(tokens.shape)
 
 # REDUCTION ####################################################################
 
 def compute_delta_activation(
     data: torch.Tensor, # (B, S, E)
-    masks: torch.Tensor, # (S,)
+    masks: torch.Tensor, # (B, S,)
     signs: torch.Tensor, # (B,)
     keepdim: bool=True,
 ) -> torch.Tensor:
     __dtype = data.dtype
     __device = data.device
+    __dim0, __dim1, __dim2 = tuple(data.shape)
     # sign each sample along the batch axis
     __shape = tuple(mlable.shapes.filter(data.shape, axes=[0]))
     __signs = signs.to(dtype=__dtype, device=__device).view(__shape)
     # combine along the batch axis to keep the shortest mask on the sequence axis
-    __shape = tuple(mlable.shapes.filter(data.shape, axes=[1]))
+    __shape = tuple(mlable.shapes.filter(data.shape, axes=[0, 1]))
     __masks = masks.to(dtype=__dtype, device=__device).view(__shape)
     # mean factor: half the signs size along the batch axis and the number of positions kept along the sequence axis
-    __factor = (0.5 * float(len(__signs)) * __masks.sum()).clamp(min=1e-8)
+    __factor = (0.5 * float(__dim0) * __masks.sum(dim=1, keepdim=True)).clamp(min=1.0)
     # take the difference along the batch axis and the average along the sequence axis
-    return (data * __signs * __masks).sum(dim=[0, 1], keepdim=keepdim) / __factor
+    return (data * __signs * __masks / __factor).sum(dim=[0, 1], keepdim=keepdim)
 
 # DELTA ########################################################################
 
@@ -147,7 +151,7 @@ def steer_model_output(
     # stop capturing activations
     __handle.remove()
     # select only the positions where the tokens differ
-    __masks = compute_sequence_mask(tokens=__inputs['input_ids'], masks=__inputs['attention_mask'])
+    __masks = compute_sequence_mask(tokens=__inputs['input_ids'])
     # activation delta at layer L
     __delta = compute_delta_activation(data=__captured[__index], masks=__masks, signs=torch.Tensor([1, -1]), keepdim=False)
     # add the delta on every forward pass

{psaiops-0.2.0.dist-info → psaiops-0.2.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: psaiops
-Version: 0.2.0
+Version: 0.2.2
 Summary: Web apps to inspect & engineer NN activations.
 License: .github/LICENSE.md
 Author: apehex

{psaiops-0.2.0.dist-info → psaiops-0.2.2.dist-info}/RECORD

@@ -5,8 +5,8 @@ psaiops/common/data.py,sha256=vGYeMN11uP9gs8rV6aSDffE_TeIX5PmdzWGwUpdGE2Y,906
 psaiops/common/dropdown.py,sha256=KA3wrfz4UlCpeoDvFAeURGK3YmVg6TBVhybaqTZa1M0,1332
 psaiops/compose/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 psaiops/compose/contrast/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-psaiops/compose/contrast/app.py,sha256=hecqh2yuUpAQMaesjIVlG-QxSJhm595Thm0CHEirZ64,8737
-psaiops/compose/contrast/lib.py,sha256=lBw-0VNvmEiY_zs9-e1TFYmOdlIbvEc7aB-9ZiZFyBU,6543
+psaiops/compose/contrast/app.py,sha256=VRg98HONjoHLfTcWLhriT2AIcB2DvM01VHlmdC-3_RA,9626
+psaiops/compose/contrast/lib.py,sha256=S5uXvnHdOA4bSbMY8Mj23KsZtIQ-zvdr4uPSUCBp-0U,6704
 psaiops/edit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 psaiops/reverse/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 psaiops/score/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -14,6 +14,6 @@ psaiops/score/attention/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG
 psaiops/score/attention/app.py,sha256=7MwFp70W0zpIcxllM2hKHN7TW3WH-BGFNK3r9a1ISr0,13444
 psaiops/score/attention/lib.py,sha256=j9p7ZXhkh5Y9VDeIe6TPUUrHRkwp1Y6P22iciaaKdL4,6932
 psaiops/steer/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-psaiops-0.2.0.dist-info/METADATA,sha256=JQ_Pu17L6AwFA3J_dxVggdzHk_iMDiLc78TDAHT9sUM,1318
-psaiops-0.2.0.dist-info/WHEEL,sha256=IYZQI976HJqqOpQU6PHkJ8fb3tMNBFjg-Cn-pwAbaFM,88
-psaiops-0.2.0.dist-info/RECORD,,
+psaiops-0.2.2.dist-info/METADATA,sha256=pZttqPtKdiI5oln701Rw-1a9mPGX8lC5oPUwabP7kYI,1318
+psaiops-0.2.2.dist-info/WHEEL,sha256=IYZQI976HJqqOpQU6PHkJ8fb3tMNBFjg-Cn-pwAbaFM,88
+psaiops-0.2.2.dist-info/RECORD,,