o
    8wi                     @   s   d dl mZ d dlmZmZ d dlmZ d dlmZ ddl	m
Z
 eG dd deZed	g d
Zee
ddddd
ee
ddddd
dZdedefddZdS )    )
namedtuple)	dataclassfield)
PeftConfig)PeftType   )llama_compute_query_statesc                       s|   e Zd ZU dZedddidZeed< edddidZe	ed< eddd	idZ
e	ed
<  fddZedefddZ  ZS )AdaptionPromptConfigz7Stores the configuration of an [`AdaptionPromptModel`].NhelpzAName of the attention submodules to insert adaption prompts into.)defaultmetadatatarget_modulesz"Number of adapter tokens to insertadapter_lenz'Number of adapter layers (from the top)adapter_layersc                    s   t    tj| _d S )N)super__post_init__r   ADAPTION_PROMPT	peft_typeself	__class__ _/home/ubuntu/sommelier/.venv/lib/python3.10/site-packages/peft/tuners/adaption_prompt/config.pyr   "   s   
z"AdaptionPromptConfig.__post_init__returnc                 C   s   dS )z1Return True if this is an adaption prompt config.Tr   r   r   r   r   is_adaption_prompt&   s   z'AdaptionPromptConfig.is_adaption_prompt)__name__
__module____qualname____doc__r   r   str__annotations__r   intr   r   propertyboolr   __classcell__r   r   r   r   r	      s   
 r	   ModelTypeConfig)compute_query_statesr   k_proj_layerv_proj_layero_proj_layer	self_attnk_projv_projo_proj)llamamistralpeft_configr   c                 C   s6   |j jtvr
tdt|j j }| jdu r|j| _| S )z1Prepare the config based on the llama model type.zHUnsupported model type for adaption prompt: '{model.config.model_type}'.N)config
model_typeTRANSFORMERS_MODEL_CONFIG
ValueErrorr   )r1   modelmodel_configr   r   r   prepare_configD   s   
r8   N)collectionsr   dataclassesr   r   peft.configr   
peft.utilsr   utilsr   r	   r&   r4   r8   r   r   r   r   <module>   s<   