o
    i+                     @   s  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	m
Z
mZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZ 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% 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/ ddl0m1Z1 ddl2m3Z3 ddl4m5Z5 ddl6m7Z7 ddl8m9Z9 ddl:m;Z; ddl<m=Z= ddl>m?Z? ddl@mAZA i de#de'de)de3ded e9d!e;d"e-d#e5d$e7d%e+d&e/d'e?d(e%d)eAd*e=d+e1d,e!iZBi dededed"edede	d ed!ed$ed%ed&ed#ed'ed(ed)ed*ed+ed,e
iZCeDeEZFG d-d. d.ZGG d/d0 d0ZHd1eIfd2d3ZJd4eIfd5d6ZKdS )7    N)OptionalUnion   )
AutoConfig)logging)
AqlmConfigAutoRoundConfig	AwqConfigBitNetQuantConfigBitsAndBytesConfigCompressedTensorsConfig
EetqConfigFbgemmFp8ConfigFineGrainedFP8Config
GPTQConfigHiggsConfig	HqqConfigQuantizationConfigMixinQuantizationMethodQuantoConfigQuarkConfig
SpQRConfigTorchAoConfig
VptqConfig   )HfQuantizer)AqlmHfQuantizer)AutoRoundQuantizer)AwqQuantizer)BitNetHfQuantizer)Bnb4BitHfQuantizer)Bnb8BitHfQuantizer)CompressedTensorsHfQuantizer)EetqHfQuantizer)FbgemmFp8HfQuantizer)FineGrainedFP8HfQuantizer)GptqHfQuantizer)HiggsHfQuantizer)HqqHfQuantizer)QuantoHfQuantizer)QuarkHfQuantizer)SpQRHfQuantizer)TorchAoHfQuantizer)VptqHfQuantizerawqbitsandbytes_4bitbitsandbytes_8bitgptqaqlmquantoquarkeetqhiggshqqzcompressed-tensors
fbgemm_fp8torchaobitnetvptqspqrfp8z
auto-roundc                   @   s.   e Zd ZdZedefddZedd ZdS )AutoQuantizationConfigz
    The Auto-HF quantization config class that takes care of automatically dispatching to the correct
    quantization config given a quantization config stored in a dictionary.
    quantization_config_dictc                 C   s   | dd }| dds| ddr"| ddrdnd}tj| }n|d u r*td|t vr>td| d	tt  t| }||S )
Nquant_methodload_in_8bitFload_in_4bit_4bit_8bitThe model's quantization config from the arguments has no `quant_method` attribute. Make sure that the model has been correctly quantizedUnknown quantization type, got  - supported types are: )	getr   BITS_AND_BYTES
ValueError AUTO_QUANTIZATION_CONFIG_MAPPINGkeyslistAUTO_QUANTIZER_MAPPING	from_dict)clsr?   r@   suffix
target_cls rS   _/home/ubuntu/maya3_transcribe/venv/lib/python3.10/site-packages/transformers/quantizers/auto.pyrO   q   s    

z AutoQuantizationConfig.from_dictc                 K   sV   t j|fi |}t|dd d u rtd| d|j}| |}|jdi | |S )Nquantization_configz)Did not found a `quantization_config` in z2. Make sure that the model is correctly quantized.rS   )r   from_pretrainedgetattrrJ   rU   rO   update)rP   pretrained_model_name_or_pathkwargsmodel_configr?   rU   rS   rS   rT   rV      s   

z&AutoQuantizationConfig.from_pretrainedN)__name__
__module____qualname____doc__classmethoddictrO   rV   rS   rS   rS   rT   r>   k   s    r>   c                   @   sd   e Zd ZdZedeeef fddZedd Z	edeeef de
e fdd	Zed
d ZdS )AutoHfQuantizerz
     The Auto-HF quantizer class that takes care of automatically instantiating to the correct
    `HfQuantizer` given the `QuantizationConfig`.
    rU   c                 K   s|   t |tr
t|}|j}|tjkr|jr|d7 }n|d7 }|t	 vr2t
d| dtt	  t| }||fi |S )NrD   rC   rF   rG   )
isinstancera   r>   rO   r@   r   rI   rA   rN   rL   rJ   rM   )rP   rU   rZ   r@   rR   rS   rS   rT   from_config   s   




zAutoHfQuantizer.from_configc                 K   s   t j|fi |}| |S )N)r>   rV   rd   )rP   rY   rZ   rU   rS   rS   rT   rV      s   
zAutoHfQuantizer.from_pretrainedquantization_config_from_argsc                 C   s   |durd}nd}t |trt |trt|}nt|}t |tttttfrK|durK|	 }|
 D ]
\}}t||| q4|dt|  d7 }|dkrTt| |S )z
        handles situations where both quantization_config from args and quantization_config from model config are present.
        NzYou passed `quantization_config` or equivalent parameters to `from_pretrained` but the model you're loading already has a `quantization_config` attribute. The `quantization_config` from the model will be used. z"However, loading attributes (e.g. z]) will be overwritten with the one you passed to `from_pretrained`. The rest will be ignored.)rc   ra   r   rO   r>   r   r	   r   r   get_loading_attributesitemssetattrrM   rL   warningswarn)rP   rU   re   warning_msgloading_attr_dictattrvalrS   rS   rT   merge_quantization_configs   s(   	



z*AutoHfQuantizer.merge_quantization_configsc                 C   s   |  dd }|  dds|  ddr"|  ddrdnd}tj| }n|d u r*td|t vrBtd| d	tt	  d
 dS dS )Nr@   rA   FrB   rC   rD   rE   rF   rG   z~. Hence, we will skip the quantization. To remove the warning, you can delete the quantization_config attribute in config.jsonT)
rH   r   rI   rJ   rK   rL   loggerwarningrM   rN   )r?   r@   rQ   rS   rS   rT   supports_quant_method   s    
z%AutoHfQuantizer.supports_quant_methodN)r\   r]   r^   r_   r`   r   r   ra   rd   rV   r   rp   staticmethodrs   rS   rS   rS   rT   rb      s    

)rb   methodc                        fdd}|S )z-Register a custom quantization configuration.c                    6    t v rtd  dt| tstd| t  < | S )NzConfig '' already registeredz*Config must extend QuantizationConfigMixin)rK   rJ   
issubclassr   rP   ru   rS   rT   register_config_fn      
z8register_quantization_config.<locals>.register_config_fnrS   )ru   r|   rS   r{   rT   register_quantization_config      
r~   namec                    rv   )zRegister a custom quantizer.c                    rw   )NzQuantizer 'rx   z!Quantizer must extend HfQuantizer)rN   rJ   ry   r   rz   r   rS   rT   register_quantizer_fn
  r}   z1register_quantizer.<locals>.register_quantizer_fnrS   )r   r   rS   r   rT   register_quantizer  r   r   )Lrj   typingr   r   models.auto.configuration_autor   utilsr   utils.quantization_configr   r   r	   r
   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   baser   quantizer_aqlmr   quantizer_auto_roundr   quantizer_awqr   quantizer_bitnetr   quantizer_bnb_4bitr    quantizer_bnb_8bitr!   quantizer_compressed_tensorsr"   quantizer_eetqr#   quantizer_fbgemm_fp8r$   quantizer_finegrained_fp8r%   quantizer_gptqr&   quantizer_higgsr'   quantizer_hqqr(   quantizer_quantor)   quantizer_quarkr*   quantizer_spqrr+   quantizer_torchaor,   quantizer_vptqr-   rN   rK   
get_loggerr\   rq   r>   rb   strr~   r   rS   rS   rS   rT   <module>   s   T	
	

)c