o
    1 i                     @   s   d dl mZmZ d dlZd dlZd dlZd dlZ	d dl
mZ d dlZd dlmZ d dlmZ d dlmZ dZdefd	d
ZdefddZdddZddefddZedkryd dlZe Zejdddd e \ZZ eej!Z"ej!s{ee"Z#dS dS dS )    )DictListN)train_test_split)tune)TuneReportCheckpointCallback)ASHASchedulerzbooster-checkpoint.jsonconfigc           	   	   C   sl   t jjdd\}}t||dd\}}}}tj||d}tj||d}tj| ||dfgdtdtd	gd
 d S )NTZ
return_X_yg      ?)Z	test_sizelabeltestF   )	frequencyfilename)Zevalsverbose_eval	callbacks)	sklearndatasetsload_breast_cancerr   xgbDMatrixtrainr   CHECKPOINT_FILENAME)	r   datalabelsZtrain_xZtest_xZtrain_yZtest_y	train_setZtest_set r   m/home/app/PaddleOCR-VL-test/.venv_paddleocr/lib/python3.10/site-packages/ray/tune/examples/xgboost_example.pytrain_breast_cancer   s   

r   c              	   C   sh   t jjdd\}}dtttt f dtttf fdd}tj||d}tj	| |ddt
|d	d
gd d S )NTr	   results_dictreturnc                 S   s   dd |   D S )Nc                 S   s   i | ]
\}}|t |qS r   )npmean).0kvr   r   r   
<dictcomp>2   s    zDtrain_breast_cancer_cv.<locals>.average_cv_folds.<locals>.<dictcomp>)items)r   r   r   r   average_cv_folds1   s   z0train_breast_cancer_cv.<locals>.average_cv_foldsr
   Fr   )Zresults_postprocessing_fnr   )r   Z
stratifiedr   )r   r   r   r   strr   floatr   r   Zcvr   )r   r   r   r(   r   r   r   r   train_breast_cancer_cv)   s   &
r+   best_resultray.tune.Resultc                 C   sB   t j| jtd}d| jd  }td| j  td|d |S )N)r         ?z
test-errorzBest model parameters: zBest model total accuracy: z.4f)r   Z	get_model
checkpointr   Zmetricsprintr   )r,   best_bstZaccuracyr   r   r   get_best_model_checkpointE   s   r2   Fuse_cvc              	   C   s   dddgt ddt g dt ddt d	d
d}tdddd}t jt j| s+tnt	ddidt j
ddd|d|d}| }| S )Nzbinary:logisticZloglosserrorr   	   )r         g      ?r.   g-C6?g?)Z	objectiveZeval_metric	max_depthZmin_child_weightZ	subsampleeta
   r6   )Zmax_tZgrace_periodZreduction_factorcpu)	resourcesztest-loglossmin)ZmetricmodeZnum_samples	scheduler)Ztune_configZparam_space)r   randintchoiceuniformZ
loguniformr   ZTunerZwith_resourcesr   r+   Z
TuneConfigfitZget_best_result)r3   Zsearch_spacer?   Ztunerresultsr   r   r   tune_xgboostO   s2   




rE   __main__z--use-cv
store_truez$Use `xgb.cv` instead of `xgb.train`.)actionhelp)r,   r-   )F)$typingr   r   numpyr!   Zsklearn.datasetsr   Zsklearn.metricsZxgboostr   Zsklearn.model_selectionr   Zrayr   Zray.tune.integration.xgboostr   Zray.tune.schedulersr   r   dictr   r+   r2   boolrE   __name__argparseArgumentParserparseradd_argumentparse_known_argsargs_r3   r,   r1   r   r   r   r   <module>   s6    

"
