o wi<@sddlmZddlmZmZddlZddlZddlZddl m Z GdddZ eGdddZ ej d d d ee fd d Zej d d d ee fddZej d d ddZej d d ddZej d d ddZej d d ddZej d d ddZdS)) dataclass)OptionalTypeN)ASRModelc@seZdZedddZdS)RNNTTestHelperNc Cst|s t|}d|vr|}|jsd|_t|}|dur2|jdg|jd}t|}n|}|durEdd|D}t|}n|}d|vrW|}|}|}|||||}t|} | d|vrptj |j dur~|j j } nd} |j | fS)NcudaTrcSsg|]}t|qS)len).0lr r [/home/ubuntu/sommelier/.venv/lib/python3.10/site-packages/tests/collections/asr/conftest.py .sz0RNNTTestHelper.wrap_and_call..)torch is_tensor FloatTensorr requires_grad LongTensorshapesumbackward synchronizegraddatacpunumpy) fnactslabelsdevice input_lengthstarget_lengthslengths label_lengthscostscostrr r r wrap_and_calls8        zRNNTTestHelper.wrap_and_call)NN)__name__ __module__ __qualname__ staticmethodr&r r r r rsrc@seZdZUeed<eed<ejed<ejed<ejed<ejed<dZeejed<dZ eejed <e dd d Z e dd dZ e dddZ e edfded dfddZe edfded dfddZdS)RnntLossSampleData vocab_sizeblank_idlogitstargetsr r!N expected_costexpected_gradsreturnc Cstgdgdgdggdgdgdggg}tddgg}dg}tgd gd gd ggd gd gdggg}tddt|tjt|tdgtdgt|tjt|dS)N)皙?333333?r3r3r3)r3r3r4r3r3)r3r3皙?g?r3)r3r3r5r3r3)gffffff?r3r5r3r3rg.ޏ@)gg6fٿ)\?r7r7)g^ʎǿ= ;Z?g lAǿr8r8)g^H ԉԿ0S ?gR4?g(?r9)d`?gd`˿r:r:r:)'a.?r;g #޿r;r;)ge)US?g"?r<r<rr,r-r.r/r r!r0r1) nparrayasarrayr+r from_numpytofloat32tensorcls activationsrr0r1r r r get_sample_smallTs<  z#RnntLossSampleData.get_sample_smallc Cstgdgdgdgdggdgdgdgdggd gd gd gd ggd gdgdgdggdgdgdgdggg}tgdg}dg}tgdgdgdgdggdgdgdgdggdgd gd!gd"ggd#gd$gd%gd&ggd'gd(gd)gd*ggg}td+d,t|tjt|td-gtd+gt|tjt|d.S)/N)?@)rJ@rL)rKrKrL)rMrLrM)rJrJrK)rJrKrK)rKrJrK)rMrMrJ)rJrM@)rJrLrN)rKrMrN)rM@rO)rJrLrO)rJrOrO)rKrLrO)rMrNrL)rMrMrK)rMrLrK)rLrMrK)rOrOrJ)rrrg`:(@)gw.g;?g$)g`gMf?g@.|?g)gV?g&V?g f)g,>g ?>g"9־)g,ȿg?ga)gJ?ga?g\)g??g;i?g@z) 6A?rPg 6Q)gg羙?gh)gT_?g@?g1d)gf?g࣡?g"V)g@?g G?g/J)ghg?g )gZg?g j{g8zo?)gw?g?gmd)gis(?gັn?gxp)g|ҿg@ ?g`?)g fk?gs˿g5 ?)gUTտg/?g ?)ੵ?rQgੵr=r6r>r?r@r+rrBrCrDrErFr r r get_sample_small_blank_last|sd $  z.RnntLossSampleData.get_sample_small_blank_lastc Csgdgdgdggdgdgdggdgdgd ggd gd gd gggd gdgdggdgdgdggdgdgdggdgdgdggg}ddg}gdgdgdggdgdgd ggd!gd"gd#ggd$gd%gd&gggd'gd(gd)ggd*gd+gd,ggd-gd.gd/ggd0gd1gd2ggg}t|}td3d4gd3d3gg}t|}td5d6t|tjt|td7d7gtd4d4gt|tjt|d8S)9N)g :?g!!ir3?g֮2?)g4{n?g`8?g82b!?)g(+0?g+"?g.4?)giڈd?ghu?gLچ?)g'F?gg7?gUa&M?)g L?ga&Q?gNxE\?)gxbu?gTw?gL?)gxU?gSY?gÏ?)gR[v&?g|0?gHyÚ?)gL'r[?gϦqܟ?g<.qy?)g ))?g0?v?g #}g?)g8 ?gb޹?g#,?)gZ%n-?g9 ~j?gGB}?)g!?g%f}X ?g`2!*H?g`Gή?g8i?)gq%3?gIś4?g@ t ?)g+?gu\l?gs'Ә?)gsA9?gS$r?g|?)gQ?gfm K?gO)g_0\Ͽgib?gE?)gtg@&g% 0?)g%_)gj#?gVsĿ)gh#`߿gZe_?g t?)gAi!܋?ggfZv?)g1 (*?g m?gJƿ)g6Xg˜Jง?g)?)gf VֿgvfVgi &Y?)g3@ϸg.G*?gh>E?)gBg/uQ?g,+\?)gÿgV`g"u@=?)g@ _Yſg~ grA\?)g_5ƿg`T?g8?)g?r?g#H^Fga_Y)?)gaκg~2@!㻿g=,X?)gK.@׿gQO ?g `H?)gȋT?g"Wgh?݄?)gJAP?gzοgqZн?)g7q(gwKW?gM@?rr6r=rr>rS)rGrHr0r1rr r r get_sample_mediums/ .   z$RnntLossSampleData.get_sample_mediumr blank_firstc Csd}|rdn|d}d}d}|rtddg}ntddg}t|d|}tjd||d|gd|d} tj|g|d} tj|g|d} t||| || | d S) NrUrrr6r=Trrrr,r-r.r/r r!) r?rArrB unsqueezerCrandrEr+) rGrWrr,r- num_framestext_lentextr/r.r r!r r r get_sample_small_random+s&z*RnntLossSampleData.get_sample_small_randomcCsd}|rdn|d}d}d}|rdnd}|r|n|d}d} tjd} | j||| |fd} t| |} tj| ||d|gd|d } tj||d ||d g|d  }tj||d|d |d g|d }t ||| | ||dS)N rrrUi)sizeTrXr6rYr= rZ) r?random RandomStaterandintrrBrCr\rElongr+)rGrWrr,r-r]r^ min_symbol max_symbol batch_sizersr_r/r.r r!r r r !get_sample_medium_random_var_sizeCs*  ""z4RnntLossSampleData.get_sample_medium_random_var_size)r2r+)r'r(r)int__annotations__rTensorr0rr1 classmethodrIrTrVrboolr`rmr r r r r+Gs&      ' > o"r+session)scoper2cCtSN)rr r r r rnnt_test_helper^rwcCrurv)r+r r r r rnn_loss_sample_datacrxrycC tdS)N%stt_en_fastconformer_transducer_largerfrom_pretrainedr r r r fast_conformer_transducer_modelh r~cCrz)Nstt_en_fastconformer_ctc_larger|r r r r fast_conformer_ctc_modelmrrcCrz)Nzparakeet-tdt_ctc-110mr|r r r r fast_conformer_hybrid_modelrrrcCrz)NzA/home/TestData/asr/canary/models/canary-1b-flash_HF_20250318.nemor restore_fromr r r r canary_1b_flashwrrcCrz)Nz;/home/TestData/asr/canary/models/canary-1b-v2_20250809.nemorr r r r canary_1b_v2|rr) dataclassesrtypingrrrr?pytestrnemo.collections.asr.modelsrrr+fixturerwryr~rrrrr r r r s2  .