&`iddlZddlmZddlmZmZmZmZddlm Z ddl Z ddl m Z ddlmZGdde ZGd d eZed Gd deZdS)N)abstractmethod)DictListOptionalUnion)Callback)TensorflowCheckpoint) PublicAPIceZdZdZgdZddeeeefffd Zde defdZ dd Z dd Z dd Z dd ZddZddZddZddZddZddZddZddZddZddZxZS) _Callbackz+Base class for Ray Train's Keras callbacks.) epoch_begin epoch_endtrain_batch_begintrain_batch_endtest_batch_begintest_batch_endpredict_batch_beginpredict_batch_end train_begin train_end test_begintest_end predict_begin predict_endvalidation_endoncttt|ts|g}t fd|Dr(t d|j|_ dS)Nc3*K|] }|jvVdSN)_allowed).0wselfs n/home/jaya/work/projects/VOICE-AGENT/VIET/agent-env/lib/python3.11/site-packages/ray/train/tensorflow/keras.py z%_Callback.__init__..%s*22!q %222222z4Invalid trigger time selected: {}. Must be one of {}) superr __init__ isinstancelistany ValueErrorformatr _on)r#r __class__s` r$r(z_Callback.__init__ s i'')))"d## B 2222r222 2 2 FMM   r&logswhenctrNotImplementedError)r#r0r1s r$_handlez_Callback._handle-s!!r&NcHd|jvr||ddSdS)Nr r.r5r#epochr0s r$on_epoch_beginz_Callback.on_epoch_begin01 DH $ $ LL} - - - - - % $r&cHd|jvr||ddSdS)Nrr7r8s r$ on_epoch_endz_Callback.on_epoch_end41 $( " " LL{ + + + + + # "r&cHd|jvr||ddSdS)Nrr7r#batchr0s r$on_train_batch_beginz_Callback.on_train_batch_begin82 $( * * LL2 3 3 3 3 3 + *r&cHd|jvr||ddSdS)Nrr7r@s r$on_train_batch_endz_Callback.on_train_batch_end<s2  ( ( LL0 1 1 1 1 1 ) (r&cHd|jvr||ddSdS)Nrr7r@s r$on_test_batch_beginz_Callback.on_test_batch_begin@s2  ) ) LL1 2 2 2 2 2 * )r&cHd|jvr||ddSdS)Nrr7r@s r$on_test_batch_endz_Callback.on_test_batch_endDs2 tx ' ' LL/ 0 0 0 0 0 ( 'r&cHd|jvr||ddSdS)Nrr7r@s r$on_predict_batch_beginz _Callback.on_predict_batch_beginHs2 DH , , LL4 5 5 5 5 5 - ,r&cHd|jvr||ddSdS)Nrr7r@s r$on_predict_batch_endz_Callback.on_predict_batch_endLrCr&cHd|jvr||ddSdS)Nrr7r#r0s r$on_train_beginz_Callback.on_train_beginPr;r&cHd|jvr||ddSdS)Nrr7rOs r$ on_train_endz_Callback.on_train_endTr>r&cHd|jvr||ddSdS)Nrr7rOs r$ on_test_beginz_Callback.on_test_beginXs1 48 # # LL| , , , , , $ #r&cHd|jvr||ddSdS)Nrr7rOs r$ on_test_endz_Callback.on_test_end\s1  ! ! LLz * * * * * " !r&cHd|jvr||ddSdS)Nrr7rOs r$on_predict_beginz_Callback.on_predict_begin`s1 dh & & LL / / / / / ' &r&cHd|jvr||ddSdS)Nrr7rOs r$on_predict_endz_Callback.on_predict_enddr;r&)rr)__name__ __module__ __qualname____doc__r rstrrr(rr5r:r=rBrErGrIrKrMrPrRrTrVrXrZ __classcell__r/s@r$r r s55H"  5d3i0      "D"""""....,,,,444422223333111166664444....,,,,----++++0000........r&r c eZdZ ddeeeefdeeeefdeeeeeeeeffffd Zdedefd Z e ded e fd Z e defd Z ded efdZxZS)RayReportCallbackrN checkpoint_onreport_metrics_onmetricsct|tr|g}t|tr|g}tt||z}t |||_||_||_dS)N)r) r)r_r*setr'r(_checkpoint_on_report_metrics_on_metrics)r#rdrerfrr/s r$r(zRayReportCallback.__init__js mS ) ) ,*OM ' - - 4!2 3  #m&7788 9 9 B)6-> r&r0returnct|jtdttt fsJ|j}nt|jtr|j|ji}nmt|jtrfd|jD}n?t|jt r%fd|jD}t|t sJ|S)Nc"i|] }|| Sro)r!metricr0s r$ z;RayReportCallback._get_reported_metrics..sQQQV QQQr&c(i|]\}}||Sroro)r!keyrpr0s r$rqz;RayReportCallback._get_reported_metrics..s0   &1c6T&\   r&)r)rktyper_r*dictitems)r#r0reported_metricss ` r$_get_reported_metricsz'RayReportCallback._get_reported_metrics|s$-$t**c4)FGGGGG = #    s + +  $ tDM/BC    t , , QQQQ4=QQQ    t , ,     59]5H5H5J5J    *D11111r& checkpointct)CSave checkpoint and report metrics corresonding to this checkpoint.r3r#rfrys r$_save_and_report_checkpointz-RayReportCallback._save_and_report_checkpoints "!r&ct)Report metrics.r3r#rfs r$_report_metricsz!RayReportCallback._report_metricss "!r&r1c0||jvs ||jvsJ||}||jv}|rLtj|j}|||tj|j ddS| |dS)NT) ignore_errors) rirjrxr from_modelmodelr}shutilrmtreepathr)r#r0r1rfshould_checkpointrys r$r5zRayReportCallback._handlest***dd6M.M.M.MM,,T22 D$77  *-8DDJ  , ,Wj A A A M*/ > > > > > >   ) ) ) ) )r&)rrN)r[r\r]rr_rrrr(rxrr r}rr5r`ras@r$rcrcis;0;3>CG   S$s)^, !d3i0 %T#YS#X >?@       $ $ 4    """)="""^" "t"""^" *D * * * * * * * * *r&rcalpha) stabilityc.eZdZdZdedefdZdefdZdS)ReportCheckpointCallbackaKeras callback for Ray Train reporting and checkpointing. .. note:: Metrics are always reported with checkpoints, even if the event isn't specified in ``report_metrics_on``. Example: .. testcode:: python ############# Using it in TrainSession ############### from ray.air.integrations.keras import ReportCheckpointCallback def train_loop_per_worker(): strategy = tf.distribute.MultiWorkerMirroredStrategy() with strategy.scope(): model = build_model() model.fit(dataset_shard, callbacks=[ReportCheckpointCallback()]) Args: metrics: Metrics to report. If this is a list, each item describes the metric key reported to Keras, and it's reported under the same name. If this is a dict, each key is the name reported and the respective value is the metric key reported to Keras. If this is None, all Keras logs are reported. report_metrics_on: When to report metrics. Must be one of the Keras event hooks (less the ``on_``), e.g. "train_start" or "predict_end". Defaults to "epoch_end". checkpoint_on: When to save checkpoints. Must be one of the Keras event hooks (less the ``on_``), e.g. "train_start" or "predict_end". Defaults to "epoch_end". rfrycHtj||dS)r{ryNraytrainreportr|s r$r}z4ReportCheckpointCallback._save_and_report_checkpoints% Z88888r&cHtj|ddS)rNrrrs r$rz(ReportCheckpointCallback._report_metricss# T22222r&N)r[r\r]r^rr r}rror&r$rrs^@99)=9999 3t333333r&r)rabcrtypingrrrrtensorflow.keras.callbacksr KerasCallbackrray.train.tensorflowr ray.util.annotationsr r rcrror&r$rs0 ............@@@@@@ 555555******Z.Z.Z.Z.Z. Z.Z.Z.z<*<*<*<*<* <*<*<*~ W)3)3)3)3)30)3)3)3)3)3r&