o
    NÆÏi  ã                   @  sl   d dl mZ d dlZd dlmZmZ d dlmZ d dlm	Z	 ddl
mZ ddlmZmZ G d	d
„ d
eƒZdS )é    )ÚannotationsN)Ú	BaseTunerÚBaseTunerLayer)Ú4TRANSFORMERS_MODELS_TO_DELORA_TARGET_MODULES_MAPPING)Úget_pattern_keyé   )ÚDeloraConfig)ÚDeloraLayerÚDeloraLinearc                      sL   e Zd ZU dZdZded< eZeZ	d‡ fd	d
„Z
dd„ Zedd„ ƒZ‡  ZS )ÚDeloraModelaR  
    Creates DeLoRA model from a pretrained transformers model.

    The method is described in detail in [TODO].

    Args:
        model ([`torch.nn.Module`]): The model to be adapted.
        config ([`DeloraConfig`]): The configuration of the DeLoRA model.
        adapter_name (`str`): The name of the adapter, defaults to `"default"`.

    Returns:
        `torch.nn.Module`: The DeLoRA model.

    **Attributes**:
        - **model** ([`~transformers.PreTrainedModel`]) -- The model to be adapted.
        - **peft_config** ([`DeloraConfig`]): The configuration of the DeLoRA model.
    Údelora_ÚstrÚprefixÚconfigr   ÚreturnÚNonec                   s   t ƒ  |¡ dS )zÌ
        A helper method to check the config when a new adapter is being added.

        Raise a ValueError if there is something wrong with the config or if it conflicts with existing adapters.

        N)ÚsuperÚ_check_new_adapter_config)Úselfr   ©Ú	__class__© úL/home/ubuntu/.local/lib/python3.10/site-packages/peft/tuners/delora/model.pyr   3   s   z%DeloraModel._check_new_adapter_configc                 K  sÀ   |d u rt dƒ‚t|j ¡ |ƒ}t|j ¡ |ƒ}	|j ||j¡}
|j |	|j¡}|
||j|j	dœ}t
|tƒrA|j|fi |¤Ž d S | j|||fi |¤Ž}|| jkrV| d¡ |  ||||¡ d S )NzCurrent Key shouldn't be `None`)ÚrÚdelora_lambdaÚmodule_dropoutÚinit_weightsF)Ú
ValueErrorr   Úrank_patternÚkeysÚlambda_patternÚgetr   r   r   r   Ú
isinstancer
   Úupdate_layerÚ_create_new_moduleÚactive_adapterÚrequires_grad_Ú_replace_module)r   Údelora_configÚadapter_nameÚtargetÚtarget_nameÚparentÚcurrent_keyÚoptional_kwargsÚr_keyÚ
lambda_keyr   r   ÚkwargsÚ
new_moduler   r   r   Ú_create_and_replace<   s"   
ü


zDeloraModel._create_and_replacec                 K  s<   t |tƒr
| ¡ }n|}t |tjjƒrt||fi |¤Ž}|S )N)r"   r   Úget_base_layerÚtorchÚnnÚLinearr
   )r(   r)   r*   r1   Útarget_base_layerr2   r   r   r   r$   _   s   

zDeloraModel._create_new_module)r   r   r   r   )Ú__name__Ú
__module__Ú__qualname__Ú__doc__r   Ú__annotations__r	   Útuner_layer_clsr   Útarget_module_mappingr   r3   Ústaticmethodr$   Ú__classcell__r   r   r   r   r      s   
 	#r   )Ú
__future__r   r5   Úpeft.tuners.tuners_utilsr   r   Ú
peft.utilsr   Úpeft.utils.otherr   r   r   Úlayerr	   r
   r   r   r   r   r   Ú<module>   s   