Added control vectors parameter and wrapper

.gitignore

@@ -0,0 +1 @@
+__pycache__

exl2_wrapper.py

@@ -0,0 +1,85 @@
+import glob
+import torch
+from gguf.gguf_reader import GGUFReader
+
+class ExLlamaV2ModuleWrapper:
+    @classmethod
+    def wrap(cls, model, vector_configs):
+        vectors = {}
+        for file in glob.glob(str(model.config.model_dir) + '-vectors/*.gguf'):
+            base = file.rsplit('-', 1)[-1].replace('.gguf', '')
+            vector, direction = base.split('__')
+            print(f"Loaded control vector: {vector}, Direction: {direction}")
+            reader = GGUFReader(file)
+
+            if reader.tensors[0].n_elements != model.config.hidden_size:
+                print(' ## Control vector n_elements ({reader.tensors[0].n_elements}) != model.config.hidden_size ({model.config.hidden_size})')
+                sys.exit()
+            layers = torch.zeros((model.config.num_hidden_layers, model.config.hidden_size), dtype=torch.float32)
+            for tensor in reader.tensors:
+                idx = int(tensor.name.split('.')[-1])
+                layers[idx] = torch.from_numpy(tensor.data.copy())
+            vectors.setdefault(vector, {})[direction] = layers
+
+        vector_configs = vector_configs.split(',')
+        control_vector = torch.zeros((model.config.num_hidden_layers, model.config.hidden_size), dtype=torch.float32)
+        for vector_config in vector_configs:
+            (vector, direction, weight) = vector_config.split(':')
+            vector_dirs = None
+            for k, v in vectors.items():
+                if vector in k:
+                    vector = k
+                    vector_dirs = v
+                    break
+            if vector_dirs is None:
+                print(' !! Error: No vector for "%s" (%s)' % (vector, vector_config))
+                continue
+
+            debias_layers = vector_dirs.get('debias', None)
+            if debias_layers is None:
+                print(' !! Error: No debias for "%s" (%s)' % (vector, vector_config))
+                continue
+
+            direction_layers = vector_dirs.get(direction, None)
+            if direction_layers is None:
+                print(' !! Error: No "%s" for "%s" (%s)' % (direction, vector, vector_config))
+                continue
+
+            try:
+                weight = float(weight)
+            except Exception as e:
+                print(' !! Non float weight %s (%s)' % (weight, vector_config))
+                weight = 1.0
+
+            print(' -- Applying %s debias and %s * %s' % (vector, direction, weight))
+            control_vector += debias_layers
+            control_vector += direction_layers * weight
+
+
+        for idx, module in enumerate(model.modules):
+            if idx == 0 or idx >= (len(model.modules) - 2) or module.name != 'MLP':
+                continue
+            model.modules[idx] = ExLlamaV2ModuleWrapper(module, control_vector)
+
+    def __init__(self, module, control_vector):
+        self.module = module
+        self.control_vector = control_vector
+
+    def __getattribute__(self, name):
+        if name == 'forward':
+            return object.__getattribute__(self, 'wrapped_forward')
+
+        try:
+            return getattr(object.__getattribute__(self, 'module'), name)
+        except AttributeError:
+            pass
+        return object.__getattribute__(self, name)
+
+    def wrapped_forward(self, *args, **kwargs):
+        x = self.module.forward(*args, **kwargs)
+        try:
+            x += self.control_vector[self.module.layer_idx].clone().to(x.device)
+        except IndexError:
+            pass
+        return x
+

test_inference.py

@@ -64,6 +64,7 @@ parser.add_argument("-sl", "--stream_layers", action = "store_true", help = "Loa
 parser.add_argument("-sp", "--standard_perplexity", choices = ["wiki2"], help = "Run standard (HF) perplexity test, stride 512 (experimental)")
 parser.add_argument("-rr", "--rank_reduce", type = str, help = "Rank-reduction for MLP layers of model, in reverse order (for experimentation)")
 parser.add_argument("-mol", "--max_output_len", type = int, help = "Set max output chunk size (incompatible with ppl tests)")
+parser.add_argument("-cv", "--control_vectors", type = str, help = "List of control vectors to apply. Format: topic:direction:weight, e.g. -cv language:simple:0.5")
 
 # Initialize model and tokenizer
 
@@ -107,6 +108,10 @@ model, tokenizer = model_init.init(
 )
 cache = None
 
+if args.control_vectors:
+    from exl2_wrapper import ExLlamaV2ModuleWrapper
+    ExLlamaV2ModuleWrapper.wrap(model, args.control_vectors)
+
 # Auto split
 
 if not model.loaded and not args.stream_layers: