88
99from transformers .models .llama .modeling_llama import LlamaDecoderLayer
1010
11- forward_old = LlamaDecoderLayer .forward
11+ forward_org = LlamaDecoderLayer .forward
1212
1313
1414def capture_and_forward (self , * args , ** kwargs ):
1515 global captured_input
1616
1717 # Prepare args tuple for TICO.convert()
1818 # Get arg_names in positional args order using inspect
19- sig = inspect .signature (forward_old )
19+ sig = inspect .signature (forward_org )
2020 args_names = [
2121 # signature includes `self`` and `kwargs``.
2222 # Just retrieve the ordinary positional inputs only
@@ -38,7 +38,7 @@ def populate_args(args_dict, filter):
3838 input_to_remove = ["use_cache" ]
3939 captured_input = populate_args (args_dict , input_to_remove )
4040
41- return forward_old (self , * args , ** kwargs )
41+ return forward_org (self , * args , ** kwargs )
4242
4343
4444# Tokenizer
@@ -82,4 +82,4 @@ def populate_args(args_dict, filter):
8282model = AutoModelForCausalLM .from_pretrained (model_name )
8383model .eval ()
8484circle_model = tico .convert (model .model .layers [0 ], captured_input )
85- circle_model .save (f"llama .decoderlayer.circle" )
85+ circle_model .save (f"tinyllama .decoderlayer.circle" )
0 commit comments