better handling of layer_list, including way to override search

Author: vgel

repeng/control.py

@@ -208,9 +208,20 @@ def model_layer_list(model: ControlModel | PreTrainedModel) -> torch.nn.ModuleLi
     if isinstance(model, ControlModel):
         model = model.model
 
-    if hasattr(model, "model"):  # mistral-like
-        return model.model.layers
-    elif hasattr(model, "transformer"):  # gpt-2-like
-        return model.transformer.h
-    else:
-        raise ValueError(f"don't know how to get layer list for {type(model)}")
+    target_suffixes = [
+        "repeng_layers",  # override
+        "model.layers",  # llama, mistral, gemma, qwen, ...
+        "transformer.h",  # gpt-2
+    ]
+    for suffix in target_suffixes:
+        candidates = [
+            v
+            for k, v in model.named_modules()
+            if k.endswith(suffix) and isinstance(v, torch.nn.ModuleList)
+        ]
+        if len(candidates) == 1:
+            return candidates[0]
+
+    raise ValueError(
+        f"don't know how to get layer list for {type(model)}! try assigning `model.repeng_layers = ...` to override this search."
+    )