
    (HJj                         d dl Z d dlmZ d dlmZmZ ddlmZ ddlm	Z	m
Z
mZmZ de j                  fdZdd	Zed
k(  r ed        e        yy)    N)Path)tree_flattentree_unflatten   )convert_to_gguf)dequantize_modelloadsaveupload_to_hubreturnc                     t        j                  d      } | j                  ddd       | j                  ddd	       | j                  d
t        dd       | j                  ddt        d        | j                  ddd       | j                  ddd       | j                  dddt               | j	                         S )Nz-Fuse fine-tuned adapters into the base model.)descriptionz--model	mlx_modelz;The path to the local model directory or Hugging Face repo.)defaulthelpz--save-pathfused_modelz!The path to save the fused model.z--adapter-pathadaptersz/Path to the trained adapter weights and config.)typer   r   z--upload-repoz-The Hugging Face repo to upload the model to.)r   r   r   z--dequantizezGenerate a dequantized model.
store_true)r   actionz--export-ggufz$Export model weights in GGUF format.z--gguf-pathzTPath to save the exported GGUF format model weights. Default is ggml-model-f16.gguf.zggml-model-f16.gguf)r   r   r   )argparseArgumentParseradd_argumentstr
parse_args)parsers    U/Users/ahmed/devFolder/claude-voice/.venv/lib/python3.12/site-packages/mlx_lm/fuse.pyparse_argumentsr      s   $$CF J  
 0  
 >	   <	   ,  
 3  
 c%	       c            	      l   t        d       t               } t        | j                  | j                  d      \  }}}|j                         D cg c]/  \  }}t        |d      r||j                  | j                        f1 }}}|r|j                  t        |             | j                  r:t        d       t        |      }|j                  dd        |j                  dd        t        | j                        }t        || j                  |||d	
       | j                   r^|d   }|dvrt#        d| d      t%        t'        |j)                                     }	t+        ||	|t-        || j.                  z               | j0                  !t3        | j                  | j0                         y y c c}}w )NzLoading pretrained modelT)adapter_pathreturn_configfuse)
dequantizezDequantizing modelquantizationquantization_configF)donate_model
model_type)llamamixtralmistralzModel type z# not supported for GGUF conversion.)printr   r	   modelr!   named_moduleshasattrr#   r$   update_modulesr   r   popr   	save_pathr
   export_gguf
ValueErrordictr   
parametersr   r   	gguf_pathupload_repor   )
argsr-   	tokenizerconfignmfused_linearsr2   r(   weightss
             r   mainr@   <   s   	
$%D#

!2!2$ E9f ''))DAq1f 
AFFdooF./)   ^M:;"# '

>4(

($/T^^$I

 L)
<<j\)LM  |E$4$4$678	7FC	DNN8R4ST#dnnd&6&67 $Es   4F0__main__zwCalling `python -m mlx_lm.fuse...` directly is deprecated. Use `mlx_lm.fuse...` or `python -m mlx_lm fuse ...` instead.)r   N)r   pathlibr   	mlx.utilsr   r   ggufr   utilsr   r	   r
   r   	Namespacer   r@   __name__r,    r   r   <module>rI      sX      2 ! *++ *Z+8\ z		H 	F r   