Zj_SSKrSSKJr SSKJr SSKJr SSKJr SSK J r SSK J r SS KJr SS KJr SS KJrJrJr SS KJr SS KJrJrJrJrJr SSKJr SSK J!r! SSK"J#r# SSK$J%r%J&r&J'r'J(r(J)r) SSK*J+r+J,r, SSK-J.r.J/r/J0r0J1r1J2r2 SSK3J4r4J5r5J6r6 SSK7J8r8 SSK9J:r:J;r;Jr>J?r?J@r@ SSKAJBrBJCrCJDrDJErE \2R"\G5rHSrI\0"SS9\"SS\8555rJ\0"SS9\"SS\555rK"S S!\@5rL"S"S#\R5rN"S$S%\R5rO"S&S'\R5rP"S(S)\%S*S+9rQ\0"S,S-\55rR"S.S/\%S*S+9rS"S0S1\'S*S+9rT\0"S2S3\(55rU"S4S5\=5rV"S6S7\R5rW"S8S9\R5rX"S:S;\R5rY"S<S=\:5rZ"S>S?\;5r[\0"S@SA\#55r\"SBSC\?5r]"SDSE\>5r^"SFSG\\\<5r_"SHSI\B5r`"SJSK\E5ra"SLSM\D5rb\0"SNSO9"SPSQ\C55rc/SRQrdg)SN)strict)nn) functional)initialization)ACT2FN)Cache)PreTrainedConfig)TorchvisionBackend) BatchFeatureget_patch_output_sizeselect_best_resolution)divide_to_patches)ChannelDimension ImageInputPILImageResamplingSizeDictget_image_size)FlashAttentionKwargs)BaseModelOutputWithPooling)PreTrainedModel) ImagesKwargsMultiModalDataProcessingKwargsProcessorMixinUnpack)PreTokenizedInput TextInput) TensorTypeTransformersKwargsauto_docstringcan_return_tuplelogging)CONFIG_MAPPING AutoConfig AutoTokenizer) LlamaConfig)LlamaAttentionLlamaDecoderLayerLlamaForCausalLMLlamaMLP LlamaModelLlamaPreTrainedModel LlamaRMSNorm)LlavaCausalLMOutputWithPastLlavaForConditionalGeneration LlavaModelLlavaModelOutputWithPastcURSnURSn[R"X4URURS9n[R "USS9n[R"S[R URS9n[R"Xv45n[URS5H.nXhn XhS-n X U n [R"XU5n XX&M0 U$)a  Compute the matrix multiplication (GEMM) for each expert sequentially. This approach is computationally inefficient, especially when dealing with a large number of experts. Args: token_states (torch.Tensor): Input tensor of shape (num_tokens, in_features). expert_weights (torch.Tensor): Weight tensor of shape (num_experts, in_features, out_features). tokens_per_expert (torch.Tensor): Number of tokens assigned to each expert. Returns: torch.Tensor: Output tensor of shape (num_tokens, out_features). rdtypedevicedim) shapetorchzerosr7r8cumsumlongcatrangematmul) token_statesexpert_weightstokens_per_expert num_tokens out_featuresoutputcumsum_num_tokens zero_tensor expert_numstartendtokensouts v/root/GenerationalWealth/GenerationalWealth/venv/lib/python3.13/site-packages/transformers/models/aria/modular_aria.pysequential_experts_gemmrRCs##A&J!''+L [[9K9KT`TgTg hF %6A>++auzz:K:R:RSK ;"BCN0034 !-Q/C(ll6*#=>u 5 Mzrhymes-ai/Aria) checkpointc\rSrSr%SrSrSrSSSSSSSS.rSr\ \ S 'S r \ \ S 'S r \ \ S 'S r \ \ S'S r\ S-\ S'Srg)AriaTextConfigba% moe_num_experts (`int`, *optional*, defaults to 8): The number of experts in the MoE layer. moe_topk (`int`, *optional*, defaults to 2): The number of top experts to route to for each token. moe_num_shared_experts (`int`, *optional*, defaults to 2): The number of shared experts. aria_text text_configcolwiserowwise)zlayers.*.self_attn.q_projzlayers.*.self_attn.k_projzlayers.*.self_attn.v_projzlayers.*.self_attn.o_projz%layers.*.mlp.shared_experts.gate_projz#layers.*.mlp.shared_experts.up_projz%layers.*.mlp.shared_experts.down_projiintermediate_sizemoe_num_expertsr$moe_topkmoe_num_shared_expertsN pad_token_id)__name__ __module__ __qualname____firstlineno____doc__ model_typebase_config_keybase_model_tp_planr\int__annotations__r^r_r`ra__static_attributes__rbrSrQrVrVbsjJ#O%.%.%.%.1:/81:"s!OSHc"#C# L#* rSrVc^\rSrSr%SrSrSS0r\\S.r Sr \ \ -S-\ S'Sr\ \-S-\ S 'S r\\\-\ S 'Sr\ S-\ S 'S r\\ S'Sr\\ S'Sr\\ S'U4SjrSrU=r$) AriaConfigze projector_patch_to_query_dict (`dict`, *optional*): Mapping of patch sizes to query dimensions. ariaimage_token_idimage_token_index)rY vision_configNrtrYr5vision_feature_layerprojector_patch_to_query_dict g{Gz?initializer_rangeFtie_word_embeddingsc >URc SSS.UlURR5VVs0sHup#[U5[U5_M snnUl[URR 55Ul[ UR[5(a;SURS'[URS"S0URD6UlO URc[S"5Ul[ UR[5(a+SUR;a[S0URD6Ul OURc[5Ul [TU]4"S0UD6 gs snnf)N)ii$idefics3_visionrhrb)rvitemsrkmaxvalues'max_value_projector_patch_to_query_dict isinstancertdictr%rYrVsuper __post_init__)selfkwargskv __class__s rQrAriaConfig.__post_init__sA  - - 52D .JNIkIkIqIqIs-tIsc!fc!fnIs-t*7:4;];];d;d;f7g4 d(($ / //@D  | ,!/0B0B<0P!Q!gTXTfTf!gD     '!/0A!B!DD  d&& - -,$BRBR2R-A0@0@AD     %-/D  ''.us!E9)rrvrYrt)rcrdrerfrgrh attribute_maprVr& sub_configsrtrr rlrYrurklistrvrsrxfloatryboolrrm __classcell__rs@rQroros J-M#1:NK48M4**T1804K&-4,.#S /.15!4$;5s#u# %%((rSroc\rSrSrSrg)AriaTextRMSNormrbNrcrdrerfrmrbrSrQrrrSrc2^\rSrSrSrU4SjrSrSrU=r$)AriaProjectorMLPz Feed-Forward Network module for the Aria Projector. Args: in_features (`int`): Input embedding dimension. hidden_features (`int`): Hidden dimension of the feed-forward network. output_dim (`int`): Output dimension. c>[TU]5 [R"XSS9Ul[R"X#SS9Ul[ SUlg)NFbiasgelu_new)r__init__rLinear linear_in linear_outract)r in_featureshidden_features output_dimrs rQrAriaProjectorMLP.__init__s> ;eL))OeL*%rSchURURU55nURU5nU$Nrrr)r hidden_statess rQforwardAriaProjectorMLP.forwards- !>?  6 rSr rcrdrerfrgrrrmrrs@rQrrs & rSrcF^\rSrSrSrSS\S\4U4SjjjrS SjrSr U=r $) AriaCrossAttentionzb Aria Cross-Attention module. Args: config (`AriaConfig`): The configuration to use. config dropout_ratec.>[TU]5 URRnURRnX@l[ R"X3SS9Ul[ R"X3SS9Ul [ R"X3SS9Ul [ R"X4SS9Ul [ R"X35Ul [ R"U5Ul[ R "U5Ul[ R "U5Ulg)NFrT) batch_first)rrrt hidden_sizenum_attention_heads num_headsrrq_projk_projv_projMultiheadAttentionmultihead_attnlinearDropoutdropout LayerNorm layer_norm layer_norm_kv)rrrrrrs rQrAriaCrossAttention.__init__s **66 ((<< "ii uE ii uE ii uE !33KX\]ii 9 zz,/ ,,{3\\+6rScURURU55nURU5nURU5nUR U5nUR XEXcS9upxUR URU55nU$)ai Forward pass of the AriaCrossAttention module. Args: key_value_states (`torch.Tensor`): Input tensor for key and value. hidden_states (`torch.Tensor`): Input tensor for query. attn_mask (`torch.Tensor`, *optional*, defaults to None): Attention mask. Returns: torch.Tensor: Output tensor after cross-attention.  attn_mask)rrrrrrrr) rkey_value_statesrrquerykeyvalue attn_output_s rQrAriaCrossAttention.forwards  DOOM:;--.>?kk*+ ,-,,U,T ll4;;{#;< rS) rrrrrrrrr)rr) rcrdrerfrgrorrrrmrrs@rQrrs*7z777"rSrcx^\rSrSrSrS\4U4SjjrS S\RS\RS-4Sjjr S r U=r $) AriaProjectoriz Aria Projector module. This module projects vision features into the language model's embedding space, enabling interaction between vision and language components. Args: config (`AriaConfig`): Configuration object for the model. rc>[TU]5 URUlURR UlURRUlURR Ul URR Ul URR Ul [R"[R "UR"UR 55Ul['U5Ul[R*"UR 5Ul[/UR URUR5Ulgr)rrrvpatch_to_query_dictrtrrrrkv_dimrYrrr Parameterr=r>rrr cross_attnrrr feed_forwardrrrs rQrAriaProjector.__init__s #)#G#G !//;;--AA**66 %11== ,,88\\%++f.\.\^b^n^n"op ,V4,,t'7'78,T-=-=t?S?SUYUdUderSNrrc.URSURSpCX@R;a*[SUSURR5S35eURUnURSUR S5R USS5nUbMURURS5nUR S5RSURS5S5nURXUS9nURURU55nU$) aH Forward pass of the Projector module. Args: key_value_states (`torch.Tensor`): Input tensor of shape (batch_size, num_patches, kv_dim). attn_mask (`torch.Tensor`, *optional*, default is None): Attention mask. Returns: `torch.Tensor`: Output tensor of shape (batch_size, query_number, output_dim). rr;zNumber of patches z: not found in patch_to_query_dict amongst possible values .Nr5r)r<rKeyErrorkeysr unsqueezerepeatrepeat_interleaverexpandsizerrr) rrr batch_size num_patches query_numqueries attention_outrPs rQrAriaProjector.forward#s0#3"8"8";=M=S=STU=VK 66 6$[M1klpmEmEmJmJmLlMMNO ,,[9 **Zi(2215<? rS) rrrrrrrrrrr) rcrdrerfrgrorr=Tensorrrmrrs@rQrrs=ff( PTATrSrcN\rSrSr%Sr\\S'\\S'\\\\S'\\S'Sr g) AriaImageProcessorKwargsiEa max_image_size (`int`, *optional*, defaults to `self.max_image_size`): Maximum image size. Must be either 490 or 980. min_image_size (`int`, *optional*, defaults to `self.min_image_size`): Minimum image size. Images smaller than this in any dimension will be scaled up. split_resolutions (`list[list[int]]`, *optional*, defaults to `self.split_resolutions`): A list of possible resolutions as (height, width) pairs for splitting high-resolution images into patches. split_image (`bool`, *optional*, defaults to `self.split_image`): Whether to split the image into patches using the best matching resolution from `split_resolutions`. max_image_sizemin_image_sizesplit_resolutions split_imagerbN) rcrdrerfrgrkrlrrrmrbrSrQrrEs) DI&rSrF)totalc^\rSrSr/SQr\r\Rr /SQr /SQr Sr Sr SrSrSrSrSrS \\4U4S jjrS \S \S \\4SjrSSS \SSS S4SjrSSS \S S4SjrSSS\\\S\SSS \S4 SjrS)S\SS\S\S\S\\\-S-S\\\-S-S\S-S\\-S-S \S!\S"\\\S-S#\SSS \ 4S$jjr!S*S%\S&\4S'jjr"S(r#U=r$$)+AriaImageProcessoriW pixel_values pixel_mask num_crops)?rrPFNTrc >URS5c'/SQnUVs/sHo3SS-USS-/PM snUS'[TU]"S0UD6 gs snf)Nr))r;r$)r;r)r;)r;)r;)r;)r;r])r$r)r$r)r$r$)r$r;)rr;)rr$)rr;)rr$)rr;)rr;)rr;)r]r;rr;rb)getrr)rrdefault_resolutionselrs rQrAriaImageProcessor.__init__gsa ::) * 2#{ Pc*dPc"qECKA+EPc*dF& ' "6"+esA original_resolutiontarget_resolutionreturncdUup4UupV[Xd- S5upx[XS- S5upXyXx-X-/$)zNGet padding size for patching, returns [left, top, right, bottom] for tvF.pad.r$)divmod) rrroriginal_heightoriginal_width target_height target_widthpaste_xr_xpaste_yr_ys rQ_get_padding_size$AriaImageProcessor._get_padding_sizemsE*='&7# l;Q? m=qA '-??rSimagez torch.Tensorresamplez7PILImageResampling | tvF.InterpolationMode | int | Nonecj[X[RS9upEURU[ XES9U5$)zFResize an image to a target resolution while maintaining aspect ratio.input_data_formatheightwidth)r rFIRSTresizer)rrrr new_height new_widths rQ_resize_for_patching'AriaImageProcessor._resize_for_patchingus8!6 8H8N8N!  {{5(*"NPXYYrSc|[X[RS9nURX25n[R "XS9$)zCPad an image to a target resolution while maintaining aspect ratio.r)padding)r rrrtvFpad)rrrnew_resolutionr!s rQ_pad_for_patching$AriaImageProcessor._pad_for_patchings6 /u[k[q[qr((Kwwu..rSgrid_pinpoints patch_sizec[U[5(d [S5e[U[R S9n[ XR5nURXU5nURXv5n[XS9n U $)a: Process an image with variable resolutions by dividing it into patches. Args: image (`torch.Tensor`): The input image to be processed (channels-first format). grid_pinpoints (`list[list[int]]`): A list of possible resolutions as (height, width) pairs. patch_size (`int`): Size of each square patch to divide the image into. resample (`PILImageResampling | tvF.InterpolationMode | int | None`): Resampling filter to use when resizing. Returns: `list[torch.Tensor]`: A list of image patches in channels-first format. z6grid_pinpoints must be a list of possible resolutions.) channel_dim)r() rr TypeErrorrrrrrr%r) rrr'r(r image_sizebest_resolution resized_image padded_imagepatchess rQget_image_patches$AriaImageProcessor.get_image_patchessk..$//TU U#E7G7M7MN 0L11%(S --mM #LHrSimages do_rescalerescale_factor do_normalize image_mean image_stddisable_groupingreturn_tensorsrrrrc XU S;a [S5e/n/nSnUGHAnU (aURUXU 5nOU/nUb[U5U:a [U5nUHnURSURSnnU [ UU5- nUU:a[ [ UU-5U 5nU nOU n[ [ UU-5U 5nUR U[UUS9U 5nU U- nU U- n[R"USSUU/5n[R"X4[RS9nSUSU2SU24'URU5 URU5 M GMD [R"USS 9nURUX#XEU5n[R"USS 9n[!UUUS .US 9$) Nrrz(max_image_size must be either 490 or 980r5rr)r7Tr9rdata tensor_type) ValueErrorr1lenr<rrkrrr"r#r=r>rappendstackrescale_and_normalizer )rr3r4r5r6r7r8r9r:rrrrrr pixel_masksprocessed_cropsrr crop_images crop_imagehwscalenew_hnew_wpadding_bottom padding_rightrstacked_images stacked_maskss rQ _preprocessAriaImageProcessor._preprocesss"  +GH H  E"44U*"":.&&z2'*<_!<33 J R[  KQ7  .+& '   rSrrcURSUR5nURSUR5n[X4UR5upgU(dSnU$Xe-U-U-nU$)aU A utility that returns number of image patches for a given image size. Args: height (`int`): Height of the input image. width (`int`): Width of the input image. images_kwargs (`dict`, *optional*): Any kwargs to override defaults of the image processor. Returns: `int`: Number of patches per image. rrr;)rrrrr) rrr images_kwargsrrresized_height resized_widthrs rQget_number_of_image_patches.AriaImageProcessor.get_number_of_image_patchessv$'' t7G7GH &**+;T=P=PQ(>PTPfPf(g%*a 1?0PS`0`dr0r rSrb)rrNFNr)%rcrdrerfmodel_input_namesr valid_kwargsrBICUBICrr7r8rrrrdo_convert_rgbr4r6rrtuplerrkrrr%r1rrstrrr rSrYrmrrs@rQrrWsC+L!))H JINNKNJL#(@!A# @U@u@Y]^aYb@ Z Z! ZL Z  Z//!/  /T#Y  L  n  V"!48!NRC ^$C C  C  C DK'$. C 4;&-C +C j(4/C C C  S ?T1C C LC !C J#crSrc8\rSrSr%Sr\\S'\\S'\\S'Srg)AriaImagesKwargsia split_image (`bool`, *optional*, defaults to `False`): Whether to split large images into multiple crops. When enabled, images exceeding the maximum size are divided into overlapping crops that are processed separately and then combined. This allows processing of very high-resolution images that exceed the model's input size limits. max_image_size (`int`, *optional*, defaults to `980`): Maximum image size (in pixels) for a single image crop. Images larger than this will be split into multiple crops when `split_image=True`, or resized if splitting is disabled. This parameter controls the maximum resolution of individual image patches processed by the model. min_image_size (`int`, *optional*): Minimum image size (in pixels) for a single image crop. Images smaller than this will be upscaled to meet the minimum requirement. If not specified, images are processed at their original size (subject to the maximum size constraint). rrrrbN) rcrdrerfrgrrlrkrmrbrSrQrbrbs rSrbcL\rSrSr%\\S'SSS.SSS.\RS.rSr g ) AriaProcessorKwargsirVF)r!return_mm_token_type_idsr)rr) text_kwargsrVr:rbN) rcrdrerfrbrlrPYTORCH _defaultsrmrbrSrQrdrds4##(- "  %,, IrSrdc ^\rSrSrSS\\-S\S-S\\\-\4S-4U4Sjjjr \ SS\ \ -\ \ -\ \ -S\S-S \\S \4S jj5rSS jr\S 5rSrU=r$) AriaProcessori-N tokenizer chat_templatesize_conversionc,>UcSSS.nUR5VVs0sHupV[U5U_M snnUlURUlURUlUbUR cUR Ul[TU]!XUS9 gs snnf)z` size_conversion (`Dict`, *optional*): A dictionary indicating size conversions for images. Nr{r|r<)rl) r~rkrm image_tokenrr pad_token unk_tokenrr)rimage_processorrkrlrmrrrs rQrAriaProcessor.__init__/s  "$'c2O6E6K6K6MN6MdaA 6MN$00'66  Y%8%8%@"+"5"5I  =Q OsBtextr3rrc bUR"[4SURR0UD6n[ U[ 5(aU/nO8[ U[ 5(d#[ US[ 5(d [S5eUbUR"U40USD6nURURRSn/nURS5U-nUHQn U RURRURRU-5n URU 5 MS O0nUnUSRS S5n USRS S 5n UR"U40USDS S0D6n UR!X|S /S 9 U (aUR#U S5U S'[%0U EUEU S9$)a[ Returns: [`BatchFeature`]: A [`BatchFeature`] with the following fields: - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`. - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not `None`). - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`. - **pixel_mask** -- Pixel mask to be fed to a model. Returned when `images` is not `None`. tokenizer_init_kwargsrzAInvalid input text. Please provide a string, or a list of stringsNrVr$rrfr:reFr) modalities input_idsmm_token_type_idsr>) _merge_kwargsrdrk init_kwargsrr`rr+rrrmrr<popreplacerorC_check_special_mm_tokenscreate_mm_token_type_idsr ) rrtr3r output_kwargs image_inputstokens_per_imageprompt_stringsrsampler:re text_inputss rQ__call__AriaProcessor.__call__Es"**  "&.."<"<  dC 6DD$'' 47C0H0H_` `  //Y-:XYL#33L4M4M4S4STU4VW N$((58HHI(B(BDNND^D^ajDjk%%f- L!N&}599:JDQ#0#?#C#CD^`e#f nn^i}]7Sidhi  %%nwi%X #/3/L/L[YdMe/fK+ ,!@K!@ L`q r   4Dm n, ,,! ss *(CCcURRnURRnUVs/sH o3S:wdM UPM nn[[R X-55$s snf)Nr)rkr[rrrrfromkeys)rtokenizer_input_namesimage_processor_input_namesnames rQr[AriaProcessor.model_input_namessb $ @ @&*&:&:&L&L#9T&k8S_jWjt8S#&kDMM"7"UVWW'ls A&A&)rorrrm)NNNNr)rcrdrerfr'r`rrrkrr!rrrrrrdr rrpropertyr[rmrrs@rQrjrj-s)-$(9= R!3&RTz R eck3./$6 RR,%)0^++d9o=EV@WW0^T!0^,- 0^  0^0^d-4XXrSrjc4^\rSrSrSrS\4U4SjjrSrU=r$)AriaSharedExpertsMLPia Shared Expert MLP for shared experts. Unlike routed experts, shared experts process all tokens without routing. This class reconfigures the intermediate size in comparison to the LlamaMLP. Args: config (`AriaTextConfig`): Configuration object for the Aria language model. rc`>[TU]U5 URUR-Ulgr)rrr\r`rs rQrAriaSharedExpertsMLP.__init__s)  !'!9!9F[TU]5 XlX lX0l[ R "[R"X1U55Ul gr) rrrrHgroupsrrr=emptyweight)rrrHrrs rQrAriaGroupedExpertsGemm.__init__s: &( ll5;;vL#QR rScL[UURUR55$)a- Perform grouped matrix multiplication. Args: input (`torch.Tensor`): Input tensor of shape (num_tokens, in_features). tokens_per_expert (`torch.Tensor`): Number of tokens assigned to each expert. Returns: torch.Tensor: Output tensor of shape (num_tokens, out_features). )rRrcpu)rinputrFs rQrAriaGroupedExpertsGemm.forwards''  KK  ! ! #  rS)rrrHrrrs@rQrrs S  rSrc\^\rSrSrS\SS4U4SjjrSrS\R4Sjr Sr U=r $) AriaExpertsirrNc>[TU]5 Xl[URUR S-UR 5Ul[UR URUR 5Ulg)Nr$) rrrrrr\r^fc1fc2rs rQrAriaExperts.__init__s_  )&*<*V>VYZ>Z\b\r\rs)&*B*BFDVDVX^XnXnorSc[R"XRRSS9up#[R R USS9nX44$)Nr;)rr:r5r9)r=topkrr_rrsoftmax)r router_logits top_logits top_indicesscoress rQroute_tokens_to_experts#AriaExperts.route_tokens_to_expertssB"'**]kk>R>RXY"Z &&zr&:""rSc$URU5up4URn[R"UR 5R [R 5URRSURRS- S9R U5nUnURS5n[R"U5n URSXRR-5n URX5n [R"U SSS9up[R R#U 5U -n UR%X5n[R&"UR(SURR-UR+S54URUR,S9nUR/SX5 URSURRUR+S55nXR1S5-R3SS9nU$)Nrr;)binsminrr5r$r9r6)rr7r=histcflattentofloat32rr^viewargsort index_selectr_rchunkrrsilurr>r<rr8 index_copy_rsum)rrr top_k_index top_k_weightsoriginal_dtyperFindicesflatten_indicessorted_indicespermuted_tokens fc1_output projectiongate expert_outputunpermuted_tokensrIs rQrAriaExperts.forwards%)%A%A-%P" $**!KK    ! $ $U]] 3,, ++a/  "^   !,,r*7'44Q++J^J^8^_XXoA  ;;z1"= ]]'' 3d: ? !KK   #dkk&:&: :M[TU]5 [R"URUR SS9Ul[U5Ul[U5Ul Xl gNFr) rrrrrr^routerrexpertsrshared_expertsrrs rQrAriaTextMoELayer.__init__sM ii 2 2F4J4JQVW "6* 26: rSrrcURnURSURS55nURU5nUR X5RU5nUR URU55nXE-$Nr5)r<rrrrr)rroriginal_shaperrshared_expert_outputs rQrAriaTextMoELayer.forward sv&,,%**2}/A/A"/EF  M2  ]BGGW #22=3E3En3UV33rS)rrrr) rcrdrerfrVrr=rrrmrrs@rQrrs/~4U\\4ell44rSrc\rSrSrSrSrg)AriaTextAttentioniz=Multi-headed attention from 'Attention Is All You Need' paperrbN)rcrdrerfrgrmrbrSrQrrsGrSrc8^\rSrSrSrS\S\4U4SjjrSrU=r $)AriaTextDecoderLayeriaG Aria Text Decoder Layer. This class defines a single decoder layer in the language model, incorporating self-attention and Mixture of Experts (MoE) feed-forward network. Args: config (`AriaTextConfig`): Configuration object for the text component of the model. layer_idx (`int`): Index of the layer. r layer_idxcD>[TU]X5 [U5Ulgr)rrrmlprrrrs rQrAriaTextDecoderLayer.__init__'s +#F+rS)r) rcrdrerfrgrVrkrrmrrs@rQrrs  ,~,#,,rSrc^\rSrSr%\\S'SrSrSS/rSr Sr Sr Sr Sr \\S .r\R$"5U4S j5rS rU=r$) AriaTextPreTrainedModeli,rmodel)rrtrrTpast_key_values)r attentionsc>[TU]U5 [U[5(a5[R "UR SURRS9 gg)Ng)meanstd) r _init_weightsrrinitnormal_rrrx)rmodulers rQr%AriaTextPreTrainedModel._init_weights=sA f% f4 5 5 LLSdkk6S6S T 6rSrb)rcrdrerfrVrlbase_model_prefixinput_modalities_no_split_modulessupports_gradient_checkpointing_skip_keys_device_placement_supports_flash_attn_supports_sdpa_supports_attention_backendrr_can_record_outputsr=no_gradrrmrrs@rQrr,sh (/1IJ&*#"3N"&-'  ]]_UUrSrcZ\rSrSr%\\S'SrSrSr\ R"5S5r Sr g) AriaPreTrainedModeliDrrFTc[R"X5 [U[5(a4[R "UR URRS9 gg)N)r) rrrrr trunc_normal_rrrx)rrs rQr!AriaPreTrainedModel._init_weightsJs@%%d3 fm , ,   v||1N1N O -rSrbN) rcrdrerfrorlr_can_compile_fullgraphrr=rrrmrbrSrQrrDs4 ""& ]]_PPrSrc0^\rSrSrS\4U4SjjrSrU=r$) AriaTextModeliQrc >[TU]U5 [R"[ UR 5Vs/sHn[ X5PM sn5UlSUlUR5 gs snf)NF) rrr ModuleListrBnum_hidden_layersrlayersgradient_checkpointing post_initrs rQrAriaTextModel.__init__Rs_  mmFKFLdLdFe fFe !& 4Fe f  ',#  gsA1)r r )rcrdrerfrVrrmrrs@rQrrQs~rSrcN^\rSrSrSS0rS\4U4Sjjr\U4Sj5rSr U=r $)AriaTextForCausalLMi[lm_head.weightzmodel.embed_tokens.weightrc>[TU]U5 [U5UlURUl[ R "URURSS9UlUR5 gr) rrrr vocab_sizerrrlm_headr rs rQrAriaTextForCausalLM.__init__^sU  "6*  ++yy!3!3V5F5FUS  rSc (>[TU]"U40UD6 gr)rr)r super_kwargsrs rQrAriaTextForCausalLM.forwardgs - -rS)rrr) rcrdrerf_tied_weights_keysrVrr!rrmrrs@rQrr[s/*,GH~..rSrc\rSrSrSrg)AriaCausalLMOutputWithPastilrbNrrbrSrQrrlrrSrc\rSrSrSrg)AriaModelOutputWithPastiprbNrrbrSrQrrprrSrc^\rSrSrS\4U4SjjrSrSS\RS\RS-S\ \ \ -S \ S-S \ \ S \\-4 S jjrSS \R"S-S\RS-S\R"S-S\R$S-S\R"S-S\S-S\RS-S\ S-S \ \S \\-4SjjrSrU=r$) AriaModelitrcD>[TU]U5 [U5Ulgr)rrrmulti_modal_projectorrs rQrAriaModel.__init__us  %26%:"rSc~UcgURSURRRURRRS9nURSURRRURRRS9nUR SS9S:R 5$)Nr;) dimensionrstepr$)r5r=r9r)unfold vision_towerrr(rr)rrpatches_subgrids rQ_create_patch_attention_mask&AriaModel._create_patch_attention_maskys  $++""))44""))44,  *00""))44""))441   ###1A5;;==rSNrrruoutput_hidden_statesrrc URU5nUR"U4USSS.UD6nSnUb'URS5n [R"U 5nUR Un UR XS9UlU$)NT)patch_attention_maskr) return_dictr;r)r'r%rr= logical_notrr pooler_output) rrrrur)rr+ image_outputsimage_attn_maskflattened_maskselected_image_features rQget_image_featuresAriaModel.get_image_featuress $@@L))  !5!%      +199!!(C!44)%%)[[%E%E 5 m  ,..}/C/C]EXEXYN!%!:!:~";" *889K\M%% )%+'    '%777@G33d!//))2>2J   QU   rS)r)Nr5N)NNNNNNNN)rcrdrerfrorr'r= FloatTensorrkrrrr r_rr3 LongTensorrr rrrrmrrs@rQrrtsa;z;>&0402,0 ''%%,"DIo  #Tk  +,  + +8.215.2.204(,26!%, ##d*, ''$., $$t+ ,  t+ , &&- , , ((4/, $;, -.,  ( (, , rSrz Aria model for conditional generation tasks. This model combines a vision tower, a multi-modal projector, and a language model to perform tasks that involve both image and text inputs. ) custom_introc^\rSrSrSS0r\SS\RS\RS-S\\ \-S\ \ S \ \ -4 S jj5r\\SS \R S-S\RS-S\R S-S \R"S-S \R S-S\S-S\RS-S\R S-S\S-S\\R"-S\ \ S \ \-4Sjj55rSU4SjjrSrU=r$)AriaForConditionalGenerationirz(model.language_model.embed_tokens.weightNrrrurrc BURR"SUUUS.UD6$)N)rrrurb)rr3)rrrrurs rQr3/AriaForConditionalGeneration.get_image_featuress3zz,, %!!5    rSrxr5r6rr7labelsr8logits_to_keepc UR"SUUUUUUUU S.U D6n U Sn [U [5(a [U *S5OU nUR U SS2USS245nSnUb3UR "SXUR RRS.U D6n[UUU RU RU RS9$)a labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` or `model.image_token_id` (where `model` is your instance of `AriaForConditionalGeneration`). Tokens with indices set to `model.image_token_id` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. Example: ```python >>> import httpx >>> from io import BytesIO >>> import torch >>> from PIL import Image >>> from io import BytesIO >>> from transformers import AutoProcessor, AutoModel >>> from transformers.image_utils import load_image >>> # Note that passing the image urls (instead of the actual pil images) to the processor is also possible >>> image1 = load_image("https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg") >>> image2 = load_image("https://cdn.britannica.com/59/94459-050-DBA42467/Skyline-Chicago.jpg") >>> image3 = load_image("https://cdn.britannica.com/68/170868-050-8DDE8263/Golden-Gate-Bridge-San-Francisco.jpg") >>> processor = AutoProcessor.from_pretrained("Rhymes-AI/Aria") >>> model = AutoModel.from_pretrained("Rhymes-AI/Aria", dtype=torch.bfloat16, device_map="auto") >>> # Create inputs >>> messages = [ ... { ... "role": "user", ... "content": [ ... {"type": "image"}, ... {"type": "text", "text": "In this image, we can see the city of New York, and more specifically the Statue of Liberty."}, ... {"type": "image"}, ... {"type": "text", "text": "What can we see in this image?"}, ... ] ... }, ... { ... "role": "user", ... "content": [ ... {"type": "image"}, ... {"type": "text", "text": "In which city is that bridge located?"}, ... ] ... } ... ] >>> prompts = [processor.apply_chat_template([message], add_generation_prompt=True) for message in messages] >>> images = [[image1, image2], [image3]] >>> inputs = processor(text=prompts, images=images, padding=True, return_tensors="pt").to(model.device) >>> # Generate >>> generated_ids = model.generate(**inputs, max_new_tokens=256) >>> generated_texts = processor.batch_decode(generated_ids, skip_special_tokens=True) >>> print(generated_texts[0]) Assistant: There are buildings, trees, lights, and water visible in this image. >>> print(generated_texts[1]) Assistant: The bridge is in San Francisco. ```)rxrrr5r6rr7r8rN)logitsrKr)lossrNrrrrb) rrrkslicer loss_functionrrYrrrrr)rrxrrr5r6rr7rKr8rLrrBr slice_indicesrNrOs rQr$AriaForConditionalGeneration.forwardsZ**  %!)%+'      8B>SV8W8W~ot4]k mA}a,?@A  %%9P9P9[9[_eD*#33!//))   rSc >[T U]"U4UUUUUS.U D6n U(dU RSS5(dXJS'XZS'U $)N)rr7r5rLis_first_iterationr8Trr)rprepare_inputs_for_generationr) rrxrr7rrr5rLrUr model_inputsrs rQrV:AriaForConditionalGeneration.prepare_inputs_for_generationXsbw<  +'))1    VZZ T%B%B ,8 ()3 &rSrbr) NNNNNNNNNr)NNNNNNF)rcrdrerfrr!r=rDrkrrr r_rr3r"rErr rrrrVrmrrs@rQrHrHs+,VW0402  ''  %%,  "DIo  +,  + +    .215.2.204(,26*.!%-.h ##d*h ''$.h $$t+ h  t+ h &&- h h ((4/h   4'h $;h ell*h +,h  + +h h Z rSrH) rorVrrjrHrrrrr)er=huggingface_hub.dataclassesrrtorchvision.transforms.v2rr"rr activationsr cache_utilsr configuration_utilsr image_processing_backendsr image_processing_utilsr r rimage_transformsr image_utilsrrrrrmodeling_flash_attention_utilsrmodeling_outputsrmodeling_utilsrprocessing_utilsrrrrrtokenization_pythonrrutilsrr r!r"r#autor%r&r'llama.configuration_llamar(llama.modeling_llamar)r*r+r,r-r.r/llava.modeling_llavar0r1r2r3 get_loggerrcloggerrRrVrorModulerrrrrrbrdrjrrrrrrrrrrrrrrH__all__rbrSrQrqs .7&! 3;aa1C:-ff?=<3   H %>+,![!-!:+,*(!*(-*(Z l ryy244n>BII>B|5$m+mm`|5* *%  lXNlXlX^ Z8 Z ) RYY) X)")))X4ryy4"HH,,,$UoUU. P. PJ.13C." !<  6 [ [ |\#@\\~ rS