Pih ddlZddlZddlmZmZmZddlmZddlm Z m Z m Z m Z m Z mZddlZddlZddlZddlmZddlmZmZmZmZddlmZdd lmZe d Z e d Z!e d Z"e d Z#de$de%fdZ&de fdZ'dedee(e)e$e*efdedejfdZ+dedee(e)e$e*efdejfdZ,dej-de%fdZ.Gdde e!e"e#fZ/de0e*e1e fde0e*e ffdZ2Gdde/ejej-ejfZ3Gdde/e0e1e0fZ4Gd d!e/e0ej5e0fZ6Gd"d#e/ej7ej8ej7fZ9Gd$d%Z:Gd&d'Z;Gd(d)eZ<Gd*d+e<Z=Gd,d-e<Z>Gd.d/e e!e"e#fZ?Gd0d1e?e!e"e#fZ@Gd2d3e?e!e"e#fZAGd4d5eAejej-ejfZBGd6d7e?ee1efZCGd8d9eAej7ej8ej7fZDGd:d;e?e0e"e0fZEde*de(ddfd?ZGdee(e)e$e*efde*fd@ZH dFdedee(e)e$e*efde edejfdAZI dGdedee(e)e$e*efdCe?dDe e1fdEZJdS)HN)IterableMappingMutableMapping)partial)AnyCallableGenericOptionalTypeVarUnion)Features)_ArrayXDExtensionType_is_zero_copy_onlydecode_nested_examplepandas_types_mapper)Table)no_op_if_value_is_nullT RowFormat ColumnFormat BatchFormatkeyreturnc8|jdko|j|jkS)N)stepstopstartrs r/home/jaya/work/projects/VOICE-AGENT/VIET/agent-env/lib/python3.11/site-packages/datasets/formatting/formatting.py_is_range_contiguousr"(s 8q= 2SX22cHtd|dt|d)NzWrong key type: 'z ' of type 'z6'. Expected one of int, slice, range, str or Iterable.) TypeErrortyper s r!_raise_bad_key_typer',s0 mCmmDIImmm  r#tableindicesc t|trY|jzddd}t ||St|tr!t| j}t|trnt|r^|j dkrSt |d|j |j |j z dDS t|trK||g}t |dSt|t rt |fd|DSt#|dS)aE Query a pyarrow Table to extract the subtable that correspond to the given key. The :obj:`indices` parameter corresponds to the indices mapping in case we cant to take into account a shuffling or an indices selection for example. The indices table must contain one column named "indices" of type uint64. rrc6g|]}|S)as_py).0is r! z5_query_table_with_indices_mapping..Cs iiia iiir#cg|]C}|dddDS)rr) fast_slicecolumnr-)r.r/r)s r!r0z5_query_table_with_indices_mapping..KsI#_#_#_VWG$6$6q!$<$<$C$CA$F$Fq$I$O$O$Q$Q#_#_#_r#N) isinstanceintr2num_rowsr3r- _query_tablesliceranger)r"rrstrselect to_pylistrr')r(rr)s `r!!_query_table_with_indices_mappingr=2s#s(  w'7!7;;BB1EEaHNNPPE3'''#u4S[[!1223#u  $ $ aii7+=+=ciTWT]I]+^+^+e+efg+h+hiii  #sB cU##E7>>!#4#4#>#>#@#@AAA#x  aE#_#_#_#_[^#_#_#_```r#c&ttr||jzdSttr!t |jtt rCtr3jdkr(|jj jz S ttr+|j fd|j DSttrjtjtjt%dkr|j ddS||jzSt)dS)zY Query a pyarrow Table to extract the subtable that correspond to the given key. rrc g|] }|k| Sr,r,)r.r3rs r!r0z _query_table..^s Z Z ZFFVYMMMMMr#N)r4r5r2r6r8r9r)r"rrr:r(drop column_namesrnpfromiterint64len fast_gatherr')r(rs `r!r7r7Psj#s9en 4a888#u2S[[001#u  $ $ a##CIsx#)/CDD D #s\{ Z Z Z Ze6H Z Z Z[[[#x  7k#rx(( s88q==;$$Q** *  u~!5666r#pa_arrayc|jdkSNr) null_count)rGs r!_is_array_with_nullsrKis   ""r#c`eZdZdZdejdefdZdejdefdZ dejde fdZ dS)BaseArrowExtractorz Arrow extractor are used to extract data from pyarrow tables. It makes it possible to extract rows, columns and batches. These three extractions types have to be implemented. pa_tablerctNNotImplementedErrorselfrNs r! extract_rowzBaseArrowExtractor.extract_rowt!!r#ctrPrQrSs r!extract_columnz!BaseArrowExtractor.extract_columnwrVr#ctrPrQrSs r! extract_batchz BaseArrowExtractor.extract_batchzrVr#N) __name__ __module__ __qualname____doc__parrrUrrXrrZr,r#r!rMrMms "BH""""""rx"L"""""bh";""""""r#rMpy_dictc>d|DS)z:Return the first element of a batch (dict) as a row (dict)c&i|]\}}||dS)rr,)r.rarrays r! z_unnest..s" < < .s-```3T// >>```r#)rArSs``r!rZz!NumpyArrowExtractor.extract_batchs$`````(J_````r#rGcZt|tjrt|jtr/t |jjdfd|jDnt |jotd|jDfd|jDnt|jtr2t |jjd| nLt |jot| |  tdkrtfdDrdtjtjd krtjt& Stjd t& Stjtjd krtjStjd S)NT)unnestcFg|]}|D]}|Szero_copy_onlyto_numpyr.chunkrowrs r!r0z=NumpyArrowExtractor._arrow_array_to_numpy..I!^l@m@m9z.sGKK8=,U333KKKKKKr#cFg|]}|D]}|Srrrs r!r0z=NumpyArrowExtractor._arrow_array_to_numpy..rr#rrc3K|]m}t|tjr&|jtkp>|jdjkp(t|t otj|VndS)rN)r4rBndarraydtypeobjectshapefloatisnan)r.xrcs r!rz.sArz**_60A0^QWPUVWPXP^E^:q%((8RXa[[r#z2.0.0b1)rF)copyrr)r4r_ ChunkedArrayr&rr storage_dtypechunksallrrKtolistrEanyrBlib NumpyVersion __version__asarrayrrc)rTrGrcrs @@r!r}z)NumpyArrowExtractor._arrow_array_to_numpysG h 0 0 X(-)>?? !3HM4OX\!]!]!]%-_"4HM!B!B"sKKAIKKKHH%-_(-)>?? X!3HM4OX\!]!]!]&//~/NN!3HM!B!B!iK_`hKiKiGi&//~/NNUUWW u::>> A 6&&r~66)CC:e6::::xE@@@@ 6  r~ . .) ; ;:e$$ $8E... .r#N)r[r\r]rzr_rrtrUrBrrXrZrmr}r,r#r!rwrws///5BH55555NrxNBJNNNNabha4aaaa$/bh$/2:$/$/$/$/$/$/r#rwczeZdZdejdejfdZdejdejfdZ dejdejfdZ dS)PandasArrowExtractorrNrc`|dtS)Nr)length types_mapper)r8 to_pandasrrSs r!rUz PandasArrowExtractor.extract_rows(~~Q~''11?R1SSSr#c|dgt|jdS)Nrr)r;rrrArSs r!rXz#PandasArrowExtractor.extract_columns7s##--;N-OOPXPefgPhiir#c8|tS)Nr)rrrSs r!rZz"PandasArrowExtractor.extract_batchs!!/B!CCCr#N) r[r\r]r_rpd DataFramerUSeriesrXrZr,r#r!rrsTBHTTTTTjrxjBIjjjjDbhD2<DDDDDDr#rc eZdZ d deedeeeeeedfffdZ dedefdZ de d ede fd Z d edefd Z dS)PythonFeaturesDecoderNfeaturestoken_per_repo_idc"||_||_dSrPrrrTrrs r!rzzPythonFeaturesDecoder.__init__s! !2r#rrcV|jr!|j||jn|SN)r)rdecode_exampler)rTrs r! decode_rowz PythonFeaturesDecoder.decode_rows.^b^ktt}++C4CY+ZZZqttr#r3 column_namecX|jr"|j|||jn|Sr)r decode_columnr)rTr3rs r!rz#PythonFeaturesDecoder.decode_columns5} DM ' ' tOe ' f f f r#batchcV|jr!|j||jn|Sr)r decode_batchrrTrs r!rz"PythonFeaturesDecoder.decode_batchs.^b^kvt}))%4CY)ZZZqvvr#rP)r[r\r]r rrtr:r boolrzrrurrr,r#r!rrsmq33 *3?GSRWX[]acgXgRhMhHi?j3333 udutuuuu D s t    w$w4wwwwwwr#rceZdZdeefdZdejdejfdZdej de dej fdZ d ejdejfd Z d S) PandasFeaturesDecoderrc||_dSrPr)rTrs r!rzzPandasFeaturesDecoder.__init__s   r#rrcjr%fdjDni}|r7|||t|<|S)Nc i|]:\}}jj||ttt|;Sr,)r_column_requires_decodingrrr)r.rfeaturerTs r!rdz4PandasFeaturesDecoder.decode_row..sT   (K=:;G 3GOPPP"r#rNrc||} t|S#t$r}td||d}~wwxYw)Nz]Custom formatting function must return a dict of sequences to be able to pick a row, but got )rrf Exceptionr%rTrNformatted_batchexcs r!rzCustomFormatter.format_rowsq++H55 ?++ +   BpBB  s& AAAc||}t|drXt|dkr2t dt |dnt d| ||jdS#t$r}t d||d}~wwxYw)NrrzTried to query a column but the custom formatting function returns too many columns. Only one column was expected but got columns .zPCustom formatting function must return a dict to be able to pick a row, but got r)rhasattrrErr%rurArLrMs r!rzCustomFormatter.format_columns++H55 ?F + + ?''))**Q..dDHI]I]I_I_D`D`ddd/ tcrtt  "8#8#;< <   tcrtt  sB$$ C.CCc||}|j|}||SrP)r"rZrrrr?s r!rzCustomFormatter.format_batchsI++--;;HEE,99%@@~~e$$$r#r )r[r\r]r^rrtrzr_rrrrrr@rAs@r!rHrHs##(D64<"8######28bh<&%RX%$%%%%%%%%r#rHcolumnsc8||vrtd|d|dS)NzColumn z5 not in the dataset. Current columns in the dataset: )KeyError)rrTs r!_check_valid_column_keyrW s5 'ddd[bddeeer#sizect|tr,|dkr ||zdks||krtd|d|dSt|trdSt|trSt |dkr>t t||t t||dSdSt|trmt |dkrXt tt||t tt||dSdSt|dS)Nrz Invalid key: z is out of bounds for size )rX) r4r5 IndexErrorr8r9rE_check_valid_index_keymaxminrr')rrXs r!r[r[%sP#s! !GGd QC4KKSSSSTSSTT T C   !  C   ! s88a<< "3s88$ 7 7 7 7 "3s88$ 7 7 7 7 7 7 < C " "! s88a<< "3s3xx==t < < < < "3s3xx==t < < < < < < < C     r#ct|tjrdSt|trdSt|tt t frdSt|dSr)r4numbersIntegralr:r8r9rr'r s r!key_to_query_typera8se#w'((u C  x C%1 2 2wr#ct|ttttt fs5 t j|}n#t$rt|YnwxYwt|trt||j n ||j n|j }t|||t||}nt|||}|S)a1 Query a Table to extract the subtable that correspond to the given key. Args: table (``datasets.table.Table``): The input Table to query from key (``Union[int, slice, range, str, Iterable]``): The key can be of different types: - an integer i: the subtable containing only the i-th row - a slice [i:j:k]: the subtable containing the rows that correspond to this slice - a range(i, j, k): the subtable containing the rows that correspond to this range - a string c: the subtable containing all the rows but only the column c - an iterable l: the subtable that is the concatenation of all the i-th rows for all i in the iterable indices (Optional ``datasets.table.Table``): If not None, it is used to re-map the given key to the table rows. The indices table must contain one column named "indices" of type uint64. This is used in case of shuffling or rows selection. Returns: ``pyarrow.Table``: the result of the query on the input table N)r))r4r5r8r9r:roperatorindexr%r'rWrAr6r[r7r=)r(rr)rX pa_subtables r! query_tablerfBs2 cCsH= > >% %.%%CC % % %  $ $ $ $ $ %#s*U%78888#*#6wENsD)))"5#.. 7sGTTT sAA! A!Frformat_columnsc<t|tr|j}n|}t|}t |j} |||S|dkr|vr |||S|||S|fd|jD}|||} |rpt| trI|fd|jD} || |} | | ntd| | S)a Format a Table depending on the key that was used and a Formatter object. Args: table (``datasets.table.Table``): The input Table to format key (``Union[int, slice, range, str, Iterable]``): Depending on the key that was used, the formatter formats the table as either a row, a column or a batch. formatter (``datasets.formatting.formatting.Formatter``): Any subclass of a Formatter such as PythonFormatter, NumpyFormatter, etc. format_columns (:obj:`List[str]`, optional): if not None, it defines the columns that will be formatted using the given formatter. Other columns are discarded (unless ``output_all_columns`` is True) output_all_columns (:obj:`bool`, defaults to False). If True, the formatted output is completed using the columns that are not in the ``format_columns`` list. For these columns, the PythonFormatter is used. Returns: A row, column or batch formatted object defined by the Formatter: - the PythonFormatter returns a dictionary for a row or a batch, and a list for a column. - the NumpyFormatter returns a dictionary for a row or a batch, and a np.array for a column. - the PandasFormatter returns a pd.DataFrame for a row or a batch, and a pd.Series for a column. - the TorchFormatter returns a dictionary for a row or a batch, and a torch.Tensor for a column. - the TFFormatter returns a dictionary for a row or a batch, and a tf.Tensor for a column. rN)rr3c3$K|] }|v|V dSrPr,r.rrgs r!rzformat_table..s/*m*m3SV^lSlSl3SlSlSlSl*m*mr#c3$K|] }|v|V dSrPr,rjs r!rzformat_table..s:@@C>????FtDFF r#rP)NF)Kr_rccollections.abcrrr functoolsrtypingrrr r r r numpyrBpandasrpyarrowr_rrfeatures.featuresrrrrr(rutils.py_utilsrrrrrr9rr"r'r5r8r:r=r7rmrKrMrtrurfrhrorrwrrrrrrr rrr&r+r0r5rCrHrWr[rarfrrr,r#r!r{sT==========DCCCCCCCCCCCCCCCuuuuuuuuuuuu333333 GCLL GK w~&& gm$$ 3e33333S  S%X=>INX<E#ueS(*J$KPRPX2#28#####"""""L+!EF""""=T#tAw,'=DaL==== -bh"(.JK$$$$$-dD$.>?$$$1/1/1/1/1/,T2:t-CD1/1/1/hDDDDD-blBIr|.STDDDwwwwwwww*&&&&&&&&@j$j$j$j$j$~j$j$j$ZLLLLLhLLL IIIIIIII $"$"$"$"$" <<=$"$"$"N"""""i < DE""" Yy, CD E E E E E^BHbh$@A E E Eiw 672nR\29bl%JK(-%-%-%-%-%ilD 89-%-%-%`fftCyfTffff !c5%&A B!#!RV!!!!&5eUC!ABs $(( ( sE5#x/ 0(e_(X ((((^&* 9 9 9 sE5#x/ 09 9 TN 9 9 9 9 9 9 r#