o i%@sddlZddlZddlmZddlmZddlmZmZm Z m Z m Z ddl Z ddl mZmZmZmZmZmZddlmZddlmZdd lmZdd lmZdd lmZdd lmZm Z m!Z!d dl"m#Z#     d4dedede ede e$de$de%de%fddZ&dee ee#fde efddZ' d5dedede e$de%de$f d d!Z(ded"eed#ed$ede)f d%d&Z*d'd(Z+d)d*Z,Gd+d,d,Z- .d6d/e e)e$fd0e$d1e$de.fd2d3Z/dS)7N)Counter)Path)CallableIterableListOptionalUnion)ConfigModel Optimizerfix_random_seedset_dropout_rateset_gpu_allocator)ConfigValidationError)Printer)Errors)ConfigSchemaPretrain)Doc) dot_to_objectload_model_from_configregistry)ExampleTFconfig output_dir resume_path epoch_resumeuse_gpusilent skip_lastc sXt|d}|dddurt|dd|dd}|dkr%|r%t|d|dd<t|} | j} tj| dtd } t | | d } td | id } | d } t | | | d |durgt |||d }nd}j d}t dd| dr|d|d| ddn|d|ddd}|jd#i|d$fdd }zyt|| dD]a}t| | | D]3\}}t|}t||}|||}|r|j|fi|| dr|| ddkr||ddq| dr|| ddks|| dd kr||n||d!_qW|s|| ddd"dSdS|s+|| ddd"ww)%Nno_printtrainingseed gpu_allocatorr initialize init_tok2vec pretraining)schemacorpusbatcher optimizer)r lossi') frequency n_save_epochz/Pre-training tok2vec layer - starting at epoch z - saving every z epoch) r2)rr5r5r5r5)widthsaligns#z# Wordsz Total LossLosszw/sFc s|rdnd}jl|rd}n d||d}|d}|dWdn1s8wYjjj|d}d d }|t |d Wdn1scwYWddSWddS1s{wYdS) Nz.tempzmodel-last.binmodelz.binwbtok2vec)nr_wordr. epoch_lossepochz log.jsonla ) use_paramsaveragesopenwriteget_refto_bytesr?r.r@srsly json_dumps)rAis_tempis_last is_temp_str save_pathfile_logr<r-rtrackerK/home/ubuntu/.local/lib/python3.10/site-packages/spacy/training/pretrain.py _save_modelFs&    "zpretrain.._save_model max_epochs n_save_everyT)rLr)rM)r8)FF)rr rrr interpolaterresolverrcreate_pretraining_model _resume_modelattrsProgressTrackerdividerrowrange enumerate ensure_docs make_updateupdater@)rrrrrr r!msg allocatornlp_configPr+r, objective row_settingsrVrAbatch_idbatchdocsr.progressrTrRrUpretrainsf           $ rrexamples_or_docsreturncCs4g}|D]}t|tr||q||jq|SN) isinstancerappend reference)rsrp eg_or_docrTrTrUrdps   rdr<cCst|d}|d||d}|}|d|Wdn1s)wY|durStdt|}|rNt | dddddd }nt t j |d ||S) Nr"zResume training tok2vec from: rbr>z model\d+\.binrrzResuming from epoch: )rinforFreadrH from_bytesresearchstrintgroup ValueErrorrE1020)r<rrr rgrP weights_data model_namerTrTrUr]zs  $ r]rpr-objective_funccCs:||\}}||j||\}}||||t|S)zPerform an update over a single batch of documents. docs (iterable): A batch of `Doc` objects. optimizer (callable): An optimizer. RETURNS loss: A float for the loss. ) begin_updateops finish_updatefloat)r<rpr-r predictionsbackpropr. gradientsrTrTrUres   recCs|jgd |Wdn1swYt||}t|jdkr6|jdkr.|jnd}||j}z |j|dgdWnt y[|d}|d }t t j j ||d w|d }||j |}|j|dgdt||d |S) aDefine a network for the pretraining. We simply add an output layer onto the tok2vec input model. The tok2vec input model needs to be a model that takes a batch of Doc objects (as a list), and returns a list of arrays. Each array in the output needs to have one row per token in the doc. The actual tok2vec layer is stored as a reference, and only this bit will be serialized to file and read back in when calling the 'train' command. )enableNTok2VecListener*r>zGive it a doc to infer shapes)X componentlayer)rrrldropout) select_pipesr'get_tok2vec_reftype__name__ upstream_nameget_piper<make_docrrE874formatvocabr )ripretrain_configr>original_tok2vecrrcreate_functionr<rTrTrUr\s(     r\cCsb|d}|durd}d}ddg|dg}t|jd||d||j}|dr/||d}|S)NrzpTo use pretrained tok2vec weights, [pretraining.component] needs to specify the component that should load them.zcomponent can't be nullr))locrg)rerrorsdescr)rrrr<rH)rirtok2vec_componentrerrrrrTrTrUrs  rc@seZdZdddZddZdS)r_@BcCs:d|_d|_d|_t|_||_t|_d|_d|_ dS)NrYr) r. prev_lossr?rwords_per_epochr/time last_time last_updater@)selfr/rTrTrU__init__s  zProgressTracker.__init__c Cs|j|7_|j|7_tdd|D}|j||7<|j|7_|j|j}||jkre|t|j}|j|_t|_|j|j }||jt |jddt |ddt |f}t |j|_ |SdS)Ncss|]}t|VqdSru)len).0docrTrTrU sz)ProgressTracker.update..r2)widthr3) r.r@sumrr?rr/rrr _smart_roundrr) rrAr.rpwords_in_batchwords_since_updatewps loss_per_wordstatusrTrTrUrfs(       zProgressTracker.updateN)r)r __module__ __qualname__rrfrTrTrTrUr_s  r_r2r4figurer max_decimalcCsRttt|}||d}|dkrtt|St||}dt|d}||S)z=Round large numbers as integers, smaller numbers as decimals.rz%.f)rrrmin)rrrn_digits n_decimal format_strrTrTrUrs   r)NNrTF)T)r2r4)0rr collectionsrpathlibrtypingrrrrrrJ thinc.apir r r r r r thinc.configrwasabirrrschemasrtokensrutilrrrexamplerrboolrrrdr]rrer\rr_rrrTrTrTrUs          "V   %