This function modifies the method of an instance of a model class. It's part from chat-GPT. It will replace the method with the new method. Currently, we only use this function to modify the attention method of a model. Do not test it further. instance:
(instance, target_class_name, target_method_name, new_method, visited_instances=None)
| 3 | import self_extend_patch as SE |
| 4 | |
| 5 | def modify_method_of_instance(instance, target_class_name, target_method_name, new_method, visited_instances=None): |
| 6 | """ |
| 7 | This function modifies the method of an instance of a model class. |
| 8 | It's part from chat-GPT. |
| 9 | It will replace the method with the new method. |
| 10 | Currently, we only use this function to modify the attention method of a model. Do not test it further. |
| 11 | |
| 12 | instance: |
| 13 | instance of a model to modify. |
| 14 | target_class_name: |
| 15 | name of the attention class to modify. E.g. 'LlamaAttention', 'GPTNeoXAttention', etc. |
| 16 | new_method: new method to replace the original method. E.g. 'self_extend_forward'. |
| 17 | It should include a parameter 'self' to be binded to the instance. |
| 18 | """ |
| 19 | target_found = False |
| 20 | if visited_instances is None: |
| 21 | visited_instances = set() |
| 22 | # Unique identifier for the instance (using id() since object's id is unique) |
| 23 | instance_id = id(instance) |
| 24 | if instance_id in visited_instances: |
| 25 | target_found = False |
| 26 | return target_found |
| 27 | # Add the instance to the already_visited set |
| 28 | visited_instances.add(instance_id) |
| 29 | |
| 30 | # Check if this instance is of the target class |
| 31 | if instance.__class__.__name__ == target_class_name: |
| 32 | bond_method = MethodType(new_method, instance) |
| 33 | setattr(instance, target_method_name, bond_method) |
| 34 | target_found = True |
| 35 | return target_found |
| 36 | elif hasattr(instance, '__dict__'): |
| 37 | for attr_name, attr_value in instance.__dict__.items(): |
| 38 | if isinstance(attr_value, object) and not isinstance(attr_value, (list, tuple, dict, set)): |
| 39 | _found = modify_method_of_instance(attr_value, target_class_name, target_method_name, new_method, visited_instances) |
| 40 | if _found: |
| 41 | target_found = True |
| 42 | elif isinstance(attr_value, (list, tuple)): |
| 43 | for item in attr_value: |
| 44 | if isinstance(item, object): |
| 45 | _found = modify_method_of_instance(item, target_class_name, target_method_name, new_method, visited_instances) |
| 46 | if _found: |
| 47 | target_found = True |
| 48 | # If attribute value is a dictionary, iterate over its values and recurse |
| 49 | # E.g, for a ModuleList, its moudels are stored in a dictionary: ._modules |
| 50 | elif isinstance(attr_value, dict): |
| 51 | for key, value in attr_value.items(): |
| 52 | if isinstance(value, object): |
| 53 | _found = modify_method_of_instance(value, target_class_name, target_method_name, new_method, visited_instances) |
| 54 | if _found: |
| 55 | target_found = True |
| 56 | # If attribute value is a set, iterate and recurse |
| 57 | elif isinstance(attr_value, set): |
| 58 | for item in attr_value: |
| 59 | if isinstance(item, object): |
| 60 | _found = modify_method_of_instance(item, target_class_name, target_method_name, new_method, visited_instances) |
| 61 | if _found: |
| 62 | target_found = True |