lens migration script updated

lens_migration.py

@@ -7,7 +7,7 @@ from logging import warn
 from pathlib import Path
 import json
 
-from tuned_lens.model_surgery import get_final_layer_norm, get_transformer_layers
+from tuned_lens.model_surgery import get_final_norm, get_transformer_layers
 from tuned_lens.load_artifacts import load_lens_artifacts
 from tuned_lens.nn import TunedLens
 from transformers.models.bloom.modeling_bloom import BloomBlock
@@ -148,7 +148,7 @@ class TunedLensOld(th.nn.Module):
 
             # Currently we convert the decoder to full precision
             self.unembedding = deepcopy(model.get_output_embeddings()).float()
-            if ln := get_final_layer_norm(model):
+            if ln := get_final_norm(model):
                 self.layer_norm = deepcopy(ln).float()
             else:
                 self.layer_norm = th.nn.Identity()
@@ -354,7 +354,7 @@ if __name__ == "__main__":
 
     tuned_lens_old = TunedLensOld.load(args.resource_id, map_location=device)
 
-    tuned_lens = TunedLens.init_from_model(
+    tuned_lens = TunedLens.from_model(
         model, bias=tuned_lens_old.config['bias'], revision=revision
     )