
    (HJj                         d dl mZ d dlmZ d dlmZ d dlmZ d dl	m
Z
mZ ddlmZ ddlmZ e G d d	e             Z G d
 dej"                        Zy)    )	dataclass)OptionalN)tree_flattentree_unflatten   )qwen2)BaseModelArgsc                   >     e Zd ZU eed<   eed<   e fd       Z xZS )	ModelArgs
model_typetext_configc                 D    d|vr | |d   |      S t         |   |      S )Nr   r   )r   r   )super	from_dict)clsparams	__class__s     `/Users/ahmed/devFolder/claude-voice/.venv/lib/python3.12/site-packages/mlx_lm/models/qwen2_vl.pyr   zModelArgs.from_dict   s-    &&"6FKKw ((    )	__name__
__module____qualname__str__annotations__dictclassmethodr   __classcell__r   s   @r   r   r      s     O) )r   r   c                   ~     e Zd Zdef fdZ	 	 ddej                  deej                     fdZd Z	e
d        Z xZS )	Modelargsc                     t         |           || _        |j                  | _        t	        j
                  t        j                  j                  |j                              | _	        y N)
r   __init__r!   r   r   r    r   r   r   language_model)selfr!   r   s     r   r$   zModel.__init__   sF    	//#kk%//*C*CDDTDT*UVr   inputsinput_embeddingsc                 *    | j                  |||      S )N)cacher(   )r%   )r&   r'   r*   r(   s       r   __call__zModel.__call__!   s$     ""%2B # 
 	
r   c                 $   t        t        |j                                     }|j                  dd        |j                  dd        t	        t        |            }i }|j                         D ]   \  }}|j                  d      sd|z   }|||<   " |S )Nvisualvision_towerzlanguage_model.)r   listitemspopr   r   
startswith)r&   weights	sanitizedkeyvalues        r   sanitizezModel.sanitize+   s     gmmo!67Hd#ND)|G,-	!--/JC>>"34'#-"IcN * r   c                 B    | j                   j                  j                  S r#   )r%   modellayers)r&   s    r   r:   zModel.layers8   s    ""((///r   )NN)r   r   r   r   r$   mxarrayr   r+   r7   propertyr:   r   r   s   @r   r    r       sV    WY W /3	

 #288,	
 0 0r   r    )dataclassesr   typingr   mlx.corecorer;   mlx.nnnn	mlx.utilsr   r    r   baser	   r   Moduler     r   r   <module>rI      sG    "    2   ) ) ) 0BII  0r   