+gdUddlZddlZddlZddlmZddlmZddlmZddlmZm Z ddl m Z m Z m Z mZmZmZmZddlZddlZddlmZdd lmZdd lmZdd lmZmZmZdd l m!Z!dd l"m#Z#m$Z$ddl%m&Z&ddl'm(Z(ddl)m*Z*ddl+m,Z,ddl-m.Z.m/Z/m0Z0m1Z1e,e2Z3dXde e4e5fdeedefdZ6dee e4e fde e4e5ffdZ7de e4e5fdee e4e ffdZ8GddeZ9GddZ:Gdde:Z;Gd d!e;Z<Gd"d#e:Z=Gd$d%e:Z>Gd&d'e:Z?Gd(d)e:Z@d*ee4fd+ZAGd,d-e:ZBGd.d/e?ZCGd0d1e:ZDGd2d3e:ZEGd4d5e:ZFGd6d7e:ZGGd8d9e:ZHd:eId;edZLGd?d@e:ZMdAZNeGdBdCZOeGdDdEZPdFZQGdGdHeZR dYdIeeRdJee&dKee(dLeSdeRf dMZT dZdOeeRdPeeeUdQeeSdJee&dKee(dRee4deRfdSZVdTeRdUeSdVeSdeRfdWZWdS)[N)Counter)deepcopy) dataclass)cycleislice)AnyCallableDictIteratorListOptionalUnion)config)DatasetInfoMixin)Features) FeatureType_align_features!_check_if_features_can_be_aligned)_reset_fsspec_lock)PythonFormatterget_format_type_from_alias) DatasetInfo) NamedSplit) table_cast) get_logger)_merge_gen_kwargs_number_of_shards_in_gen_kwargs_shuffle_gen_kwargs_split_gen_kwargsbatch try_featuresreturnctj|}|P t|tj|j}n'#t tjtjf$rYnwxYwtj |jSN) paTable from_pydictrschematype TypeError ArrowInvalidArrowNotImplementedErrorrfrom_arrow_schema)r!r"pa_tables 9lib/python3.11/site-packages/datasets/iterable_dataset.py_infer_features_from_batchr1sx##E**H !(BIl6G,H,HIIHH2?B,GH    D   %ho 6 66s'A !A/.A/examplescrdD}fd|D}tt||S)Nci|] }|D]}|d Sr%.0examplecols r0 z&_examples_to_batch..*s' A A A' A A#C A A A Ac.g|]fdDS)c:g|]}|Sr5)getr6s r0 z1_examples_to_batch...,s%888Gw{{3888r;r5)r7r9r2s @r0r?z&_examples_to_batch..,s/ I I IS8888x888 I I Ir;)dictzip)r2colsarrayss` r0_examples_to_batchrD'sJ B AX A A AD I I I ID I I IF D&!! " ""r;c#Kt|tt|}t|D]$fd|DV%dS)z3Convert a batch (dict of examples) to examples listc(i|]\}}||Sr5r5)r7r9arrayis r0r:z&_batch_to_examples..4s#===esE!H===r;N)lennextiterrangeitems)r! n_examplesrHs @r0_batch_to_examplesrO0ssU4U ,,-..J :  >>====u{{}}=======>>r;c*eZdZdZdZdZdZdZdS)HasNextIteratorzyIterator with an hasnext() function. Taken from https://stackoverflow.com/questions/1966591/has-next-in-python-iterators.c<t||_d|_dSr%)rKit_hasnext)selfrSs r0__init__zHasNextIterator.__init__:sr(( r;c|Sr%r5rUs r0__iter__zHasNextIterator.__iter__>s r;cZ|jr|j}nt|j}d|_|Sr%)rT_thenextrJrS)rUresults r0__next__zHasNextIterator.__next__As. = #]FF$']]F  r;c|j9 t|j|_d|_n#t$r d|_YnwxYw|jS)NTF)rTrJrSr[ StopIterationrXs r0hasnextzHasNextIterator.hasnextIs[ =  % $TW  !% ! & & & %  &}s *>>N)__name__ __module__ __qualname____doc__rVrYr]r`r5r;r0rQrQ7s\DDr;rQceZdZdZdZdejjddfdZde de ddfdZ de de de e fd Z e de fd Zd S) _BaseExamplesIterablez?Base class for the examples iterable used by an IterableDatasetc@tt|d)zWAn examples iterable should yield tuples (example_key, example) of type (int/str, dict)z doesn't implement __iter__ yetNotImplementedErrorr*rXs r0rYz_BaseExamplesIterable.__iter__W!T$ZZ"P"P"PQQQr; generatorr#c@tt|d)z Either shuffle the shards/sources of the dataset, or propagate the shuffling to the underlying iterable. If the order of the shards must stay fixed (when using .skip or .take for example), then this method returns self. z+ doesn't implement shuffle_data_sources yetrhrUrks r0shuffle_data_sourcesz*_BaseExamplesIterable.shuffle_data_sources[s! "T$ZZ"\"\"\]]]r; worker_id num_workersc@tt|d)ZEither keep only the requested shard, or propagate the request to the underlying iterable.z) doesn't implement shard_data_sources yetrhrUrorps r0shard_data_sourcesz(_BaseExamplesIterable.shard_data_sourcesbs!T$ZZ"Z"Z"Z[[[r;cHtt||j|Sr%)listrLn_shardsrss r0split_shard_indices_by_workerz3_BaseExamplesIterable.split_shard_indices_by_workerfsE)T]K@@AAAr;c@tt|d)Nz doesn't implement n_shards yetrhrXs r0rwz_BaseExamplesIterable.n_shardsirjr;N)rarbrcrdrYnprandom Generatorrnintrtr rxpropertyrwr5r;r0rfrfTsIIRRR^bi.A^F]^^^^\C\c\F]\\\\BsBBQUVYQZBBBBR#RRRXRRRr;rfcxeZdZdedefdZdZdejj ddfdZ de d e ddfd Z e de fd Zd S) ExamplesIterablegenerate_examples_fnkwargsc"||_||_dSr%rr)rUrrs r0rVzExamplesIterable.__init__os$8! r;c#:K|jdi|jEd{VdS)Nr5rrXs r0rYzExamplesIterable.__iter__ss9,4,;;t{;;;;;;;;;;;r;rkr#c8t|j|j|Sr%)#ShuffledDataSourcesExamplesIterablerrrms r0rnz%ExamplesIterable.shuffle_data_sourcesvs243Ldk[deeer;rorpct|j|j|||}t fd|D}t |j|S)Keep only the requested shard.) max_num_jobsc g|] }| Sr5r5)r7rHgen_kwargs_lists r0r?z7ExamplesIterable.shard_data_sources..}s1\1\1\/!2D1\1\1\r;)r rrwrxrrr)rUrorp shard_indicesrequested_gen_kwargsrs @r0rtz#ExamplesIterable.shard_data_sourcesyse+DKdmTTT::9kRR 01\1\1\1\m1\1\1\]] 9;OPPPr;c*t|jSr%)rrrXs r0rwzExamplesIterable.n_shardss.t{;;;r;N)rarbrcr r@rVrYrzr{r|rnr}rtr~rwr5r;r0rrnsXt<<<fbi.AfFXffffQCQcQFXQQQQ<#<<<X<<.s===a3q6===r;r)rUidxrs @r0rYzSelectColumnsIterable.__iter__sS( > >HC====4+<==== = = = = > >r;rkr#c\t|j||jSr%)rrrnrrms r0rnz*SelectColumnsIterable.shuffle_data_sourcess'$T%5%J%J9%U%UW[Whiiir;rorpc^t|j|||jSr%)rrrtrrss r0rtz(SelectColumnsIterable.shard_data_sourcess*$T%5%H%HT_%`%`bfbstttr;c|jjSr%rrwrXs r0rwzSelectColumnsIterable.n_shards((r;N)rarbrcrfr strrVrYrzr{r|rnr}rtr~rwr5r;r0rrs)$9)c))))>>>jbi.AjF]jjjjuCucuF]uuuu)#)))X)))r;rc|eZdZdededefdZdZdejj ddfdZ d ed eddfd Z e defd Z d S)StepExamplesIterablerstepoffsetc0||_||_||_dSr%)rrr)rUrrrs r0rVzStepExamplesIterable.__init__s&  r;c#Kt|j} tt||j}t ||jkr||jVndSMr%)rKrrvrrrIr)rU ex_iteratorr!s r0rYzStepExamplesIterable.__iter__sg4+,,   TY7788E5zzDK''DK(((((  r;rkr#cjt|j||j|jSNrr)rrrnrrrms r0rnz)StepExamplesIterable.shuffle_data_sourcess7#   1 1) < <49UYU`    r;rorpclt|j|||j|jSr)rrrtrrrss r0rtz'StepExamplesIterable.shard_data_sourcess9#   / / ; G Gdi`d`k    r;c|jjSr%rrXs r0rwzStepExamplesIterable.n_shardsrr;N)rarbrcrfr}rVrYrzr{r|rnrtr~rwr5r;r0rrs$9c  bi.A F\     C c F\    )#)))X)))r;rceZdZ ddeedeefdZdZdZ de j j ddfd Z edefd Zd ed eddfd ZdS)#CyclingMultiSourcesExamplesIterablefirst_exhausted ex_iterablesstopping_strategych||_||_|dkr tjn tj|_dS)N all_exhausted)rrrzallanybool_strategy_func)rUrrs r0rVz,CyclingMultiSourcesExamplesIterable.__init__s8)!2.?/-Q-Q"&&XZX^r;c^ttt|jSr%)rrLrIrrXs r0_give_indice_iteratorz9CyclingMultiSourcesExamplesIterable._give_indice_iterators#U3t01122333r;c#Kd|jD}|}tjt |jd}|D]} t ||V||s:d||<||rdSt|j|||<o#t$r!d||<||rYdSYwxYwdS)Nc,g|]}t|Sr5)rQr7rs r0r?z@CyclingMultiSourcesExamplesIterable.__iter__..s WWWk_[11WWWr;FT) rrrzfullrIrJr`rrQr_)rU iteratorsindices_iterator is_exhaustedrHs r0rYz,CyclingMultiSourcesExamplesIterable.__iter__s*WWTEVWWW 5577ws4#455u== !  A 9Q<(((((!|++--I&*LO..|<<#243DQ3G#H#HIaL    #' Q**<88EEE    sA C"C$C+*C+rkr#cTfd|jD}t||jS)z*Shuffle each underlying examples iterable.c:g|]}|Sr5rnr7rrks r0r?zLCyclingMultiSourcesExamplesIterable.shuffle_data_sources..'iii 88CCiiir;)rrrrUrkrs ` r0rnz8CyclingMultiSourcesExamplesIterable.shuffle_data_sourcess3iiiiW[Whiii 2<AWXXXr;c>td|jDS)Nc3$K|] }|jV dSr%rwrs r0 z?CyclingMultiSourcesExamplesIterable.n_shards..%MMK;'MMMMMMr;minrrXs r0rwz,CyclingMultiSourcesExamplesIterable.n_shards"MM4;LMMMMMMr;rorpcVtfd|jD|jS)rrc<g|]}|Sr5rtr7iterablerpros r0r?zJCyclingMultiSourcesExamplesIterable.shard_data_sources..) c c cXX ( (K @ @ c c cr;r)rrrrss ``r0rtz6CyclingMultiSourcesExamplesIterable.shard_data_sourcess?2 c c c c cQUQb c c c"4    r;N)r)rarbrcr rfr rrVrrYrzr{r|rnr~r}rwrtr5r;r0rrs\m__ !67_LTUXM____444:Ybi.AYFkYYYY N#NNNXN C c Fk      r;rceZdZdZdeefdZdZdej j ddfdZ e de fdZd e d e ddfd Zd S) 2VerticallyConcatenatedMultiSourcesExamplesIterablea VerticallyConcatenatedMultiSourcesExamplesIterable simply chains the input iterables. It doesn't require the examples iterables to always yield the same columns. Instead, this is handled by the `IterableDataset` class or `TypedExamplesIterable`. For information, `IterableDataset` merges the features of all the datasets to concatenate into one. We use `IterableDataset._resolve_features` to obtain the features of all the datasets to concatenate. Then for each example, `IterableDataset` and `TypedExamplesIterable` automatically fill missing columns with None. This is done with `_apply_feature_types_on_example`. rc||_dSr%rrUrs r0rVz;VerticallyConcatenatedMultiSourcesExamplesIterable.__init__(r;c#.K|jD] }|Ed{V dSr%r)rUrs r0rYz;VerticallyConcatenatedMultiSourcesExamplesIterable.__iter__s=, # #K" " " " " " " " " # #r;rkr#ct}t|j}||fd|D}t |S)zTShuffle the list of examples iterable, as well as each underlying examples iterable.c:g|]}|Sr5rrs r0r?z[VerticallyConcatenatedMultiSourcesExamplesIterable.shuffle_data_sources..#s'ddd 88CCdddr;)rrvrshuffler)rUrkrrs ` r0rnzGVerticallyConcatenatedMultiSourcesExamplesIterable.shuffle_data_sourcess\y!!D-..  L!!!ddddWcddd A,OOOr;c>td|jDS)Nc3$K|] }|jV dSr%rrs r0rzNVerticallyConcatenatedMultiSourcesExamplesIterable.n_shards..(rr;rrXs r0rwz;VerticallyConcatenatedMultiSourcesExamplesIterable.n_shards&rr;rorpcHtfd|jDS)rrc<g|]}|Sr5rrs r0r?zYVerticallyConcatenatedMultiSourcesExamplesIterable.shard_data_sources../rr;)rrrss ``r0rtzEVerticallyConcatenatedMultiSourcesExamplesIterable.shard_data_sources*s7B c c c c cQUQb c c c   r;Nrarbrcrdr rfrVrYrzr{r|rnr~r}rwrtr5r;r0rrs  )T*?%@))))###P,P =PPPPN#NNNXN  +. =      r;rrct|tdDs!fdD}td|ddS)zBCheck the column names to make sure they don't contain duplicates.c3"K|] }|dkV dS)rNr5)r7counts r0rz&_check_column_names..6s&88euz888888r;c,g|]}|dk|S)rr5)r7r9counters r0r?z'_check_column_names..7s'IIIc q8H8Hc8H8H8Hr;zAThe examples iterables can't have duplicated columns but columns z are duplicated.N)rrvalues ValueError)rduplicated_columnsrs @r0_check_column_namesr3s~l##G 88w~~'7'7888 8 8 IIIIWIII tPb t t t     r;ceZdZdZdeefdZdZdej j ddfdZ e de fdZd e d e ddfd Zd S) 4HorizontallyConcatenatedMultiSourcesExamplesIterablea- HorizontallyConcatenatedMultiSourcesExamplesIterable merges examples together for the input list of iterables. It also checks that there are no duplicate columns (otherwise we don't know which one to keep). This check is done once when yielding the first example. However it doesn't fill missing columns with None. Instead, this is handled by the `IterableDataset` class or `TypedExamplesIterable`. For information, `IterableDataset` merges the features of all the datasets to concatenate into one. We use `IterableDataset._resolve_features` to obtain the features of all the datasets to concatenate. Then for each example, `IterableDataset` and `TypedExamplesIterable` automatically fill missing columns with None. This is done with `_apply_feature_types_on_example`. rc||_dSr%rrs r0rVz=HorizontallyConcatenatedMultiSourcesExamplesIterable.__init__Mrr;c#Kd|jD}tjD]}g}g}t|D]d} t |\}}||||@#t $r||YawxYw|ra|dkrtd|Di}|D]}| |d d|D} | |fVdSdS)Nc,g|]}t|Sr5)rKrs r0r?zQHorizontallyConcatenatedMultiSourcesExamplesIterable.__iter__..Qs OOOk[))OOOr;rcg|] }|D]}| Sr5r5)r7r8 column_names r0r?zQHorizontallyConcatenatedMultiSourcesExamplesIterable.__iter__..^s'(h(h(h`g(h(hQ\(h(h(h(hr;_c34K|]}t|VdSr%rr7keys r0rzPHorizontallyConcatenatedMultiSourcesExamplesIterable.__iter__..b("<"<3s88"<"<"<"<"<".vrr;)rrrss ``r0rtzGHorizontallyConcatenatedMultiSourcesExamplesIterable.shard_data_sourcesqs7D c c c c cQUQb c c c   r;Nrr5r;r0rr=s  )T*?%@))))., ? #X  +. ?      r;rc eZdZ ddejjdeeedee ffd Z e ddejjd e d eeed e e fd Zd Zdejjd dfdZde de d dfdZxZS)+RandomlyCyclingMultiSourcesExamplesIterableNrrk probabilitiesrct||t||_||_dSr%)rrVrrkr)rUrrkrrrs r0rVz4RandomlyCyclingMultiSourcesExamplesIterable.__init__{s= '8999!),,*r;r num_sourcespr#c#K|* d|d||DEd{V) d||||DEd{V))z]Get an infinite iterator that randomly samples the index of the source to pick examples from.NTc34K|]}t|VdSr%r}r7rHs r0rzSRandomlyCyclingMultiSourcesExamplesIterable._iter_random_indices..(aaqCFFaaaaaar;rsizec34K|]}t|VdSr%rrs r0rzSRandomlyCyclingMultiSourcesExamplesIterable._iter_random_indices..rr;)rr)integerschoice)rrrandom_batch_sizers r0_iter_random_indicesz@RandomlyCyclingMultiSourcesExamplesIterable._iter_random_indicess 9 baaCLLKN_L,`,`aaaaaaaaaa b baaCJJ{IZ^_J,`,`aaaaaaaaaa br;ct|j}||t|j|jS)N)r)rrkr$rIrr)rUrs r0rzARandomlyCyclingMultiSourcesExamplesIterable._give_indice_iterators:t~&&((c$2C.D.DHZ([[[r;cXfd|jD}t||jS)z;Shuffle the data sources of each wrapped examples iterable.c:g|]}|Sr5rrs r0r?zTRandomlyCyclingMultiSourcesExamplesIterable.shuffle_data_sources..rr;)rkr)rrrrs ` r0rnz@RandomlyCyclingMultiSourcesExamplesIterable.shuffle_data_sourcessAiiiiW[Whiii : IT=O    r;rorpcltfd|jD|j|j|jS)rrc<g|]}|Sr5rrs r0r?zRRandomlyCyclingMultiSourcesExamplesIterable.shard_data_sources..rr;)rrrkrrrss ``r0rtz>RandomlyCyclingMultiSourcesExamplesIterable.shard_data_sourcessE: c c c c cQUQb c c c N    "    r;)Nr)rN)rarbrcrzr{r|r r floatrrV staticmethodr}r r$rrnrtrrs@r0rrzsJ 04+< + +9& + U , + $C= + + + + + +#' b b Y  b b DK b # b b b\ b\\\  bi.A Fs     C c Fs        r;rceZdZ ddedededeeeded ee d ed eeed ee fd Z dZ de jjddfdZde de ddfdZede fdZdS)MappedExamplesIterableFNrrfunction with_indices input_columnsbatched batch_sizedrop_last_batchremove_columns fn_kwargsc ||_||_||_||_||_||_||_||_| pi|_dSr%) rr.r1r2r3r4r/r0r5) rUrr.r/r0r1r2r3r4r5s r0rVzMappedExamplesIterable.__init__sP'   $.,(*"br;c#Kt|j}d|jr|D]\}}|j |jdkr|nt ||jdz }||fgd|Dz}t |\}}|jr-|j&|jdkrt||jkrdSt|}||j gnfd|j D} |j r;| fdtt|Dt||j| i|j|jr |jD]} | =rht%tfdD} | r;t'd| dfd | Dd d td d d|D} t+t-D] \} }| |fV | dzz dS|D]\}}t|}||j gnfd|j D} |j r| t|}||j| i|j|jr |jD]} || =||fVdz dS)Nrrcg|] \}}||f Sr5r5r7rr8s r0r?z3MappedExamplesIterable.__iter__..!7j7j7j<3PWg7j7j7jr;c g|] }| Sr5r5r7r9inputss r0r?z3MappedExamplesIterable.__iter__..LwLwLw]`VTW[LwLwLwr;cg|]}|zSr5r5r7rH current_idxs r0r?z3MappedExamplesIterable.__iter__..)a)a)aa+/)a)a)ar;clg|]0}t|tk.|1Sr5rI)r7r9 first_coltransformed_batchs r0r?z3MappedExamplesIterable.__iter__..sK   0566#>OPY>Z:[:[[[[[;z!Column lengths mismatch: columns z have length c:g|]}t|Sr5rD)r7r9rFs r0r?z3MappedExamplesIterable.__iter__..sBXIXIXIx{X[\mnq\rXsXsXIXIXIr;z while z has length .rc34K|]}t|VdSr%rrs r0rz2MappedExamplesIterable.__iter__..rr;c g|] }| Sr5r5r<s r0r?z3MappedExamplesIterable.__iter__..r>r;)rKrr1r2rrAr3rIrDr0r/rrLr@rr.r5r4rJrr enumeraterO)rUiteratorrr8iterator_batchkey_examples_listr r2r! function_argsrbad_colsr  batch_idxtransformed_examplerArEr=rFs @@@@r0rYzMappedExamplesIterable.__iter__s9()) <? ! (, -, - W.$/Q2F2FH$/A*=>> '*7^$47j7j[i7j7j7j$j!!$&7!8h(3!++H 77FF*844,0,>,FLwLwLwLwdhdvLwLwLw $c!(()a)a)a)a5M^I_I_C`C`)a)a)abbb$(KK!!(( )X)X)XYYY&1!011-a00$  $T*;%<%< = =I     #4   H  (MMMXIXIXIXI@HXIXIXIMMR[MMilm~HmIiJiJMMM(("<",FLwLwLwLwdhdvLwLwLw $6!((555&*7mm##**=4=-+Z4>+Z+Z[[[&3!033/22.....q ! ! !r;rkr#c t|j||j|j|j|j|j|j|j S)&Shuffle the wrapped examples iterable.r.r/r0r1r2r4r5) r-rrnr.r/r0r1r2r4r5rms r0rnz+MappedExamplesIterable.shuffle_data_sourcessS%   1 1) < <]*,L.n    r;rorpc t|j|||j|j|j|j|j|j|j S)rrU) r-rrtr.r/r0r1r2r4r5rss r0rtz)MappedExamplesIterable.shard_data_sourcessU%   / / ; G G]*,L.n    r;c|jjSr%rrXs r0rwzMappedExamplesIterable.n_shards rr;)FNFrFNN)rarbrcrfr boolr r rr}r@rVrYrzr{r|rnrtr~rwr5r;r0r-r-sB #-1$( %.2$())*)) )  S * )  )SM))!c+)D>)))),B!B!B!H  bi.A  F^       C  c  F^     )#)))X)))r;r-ceZdZ ddedededeeeded ee f d Z d Z d ee d dfdZ de de d dfdZ ed e fdZdS)FilteredExamplesIterableFNrrr.r/r0r1r2cZ||_||_||_||_||_||_dSr%)rr.r1r2r/r0)rUrr.r/r0r1r2s r0rVz!FilteredExamplesIterable.__init__&s7'   $(*r;c #Kt|j}d|jr|D]\}}|j |jdkr|nt ||jdz }||fgd|Dz}t |\}}t |}||jgnfd|jD} |jr;| fdtt|D|j | } tt || D]\} \} } | r| V| dzz dS|D]d\}}t||jgnfd|jD} |jr| |j | } | r||fVdz edS)Nrrcg|] \}}||f Sr5r5r9s r0r?z5FilteredExamplesIterable.__iter__..Ar:r;c g|] }| Sr5r5r<s r0r?z5FilteredExamplesIterable.__iter__..Fr>r;cg|]}|zSr5r5r@s r0r?z5FilteredExamplesIterable.__iter__..HrBr;c g|] }| Sr5r5r<s r0r?z5FilteredExamplesIterable.__iter__..Sr>r;)rKrr1r2rrArDr0r/rrLrIr.rKr@)rUrLrr8rMrNr r2r!rOmaskrQ key_exampleto_keeprAr=s @@r0rYz!FilteredExamplesIterable.__iter__6s8()) < ! ( - - W.$/Q2F2FH$/A*=>> '*7^$47j7j[i7j7j7j$j!!$&7!8h*844,0,>,FLwLwLwLwdhdvLwLwLw $c!(()a)a)a)a5M^I_I_C`C`)a)a)abbb$t}m49B3GXZ^C_C_9`9`**5I5 W*))))y1}, ) - -,!) ! ! Wg,0,>,FLwLwLwLwdhdvLwLwLw $6!((555'$-7'w,&&&q  ! !r;seedr#ct|j||j|j|j|j|jS)rTr.r/r0r1r2)rZrrnr.r/r0r1r2)rUrds r0rnz-FilteredExamplesIterable.shuffle_data_sources[sH'   1 1$ 7 7]*,L     r;rorpct|j|||j|j|j|j|jS)rrf)rZrrtr.r/r0r1r2rss r0rtz+FilteredExamplesIterable.shard_data_sourcesfsJ'   / / ; G G]*,L     r;c|jjSr%rrXs r0rwz!FilteredExamplesIterable.n_shardsqrr;)FNFr)rarbrcrfr rXr r rr}rVrYrnrtr~rwr5r;r0rZrZ%s  #-1$(++*++ +  S * +  +SM++++ #!#!#!J  #  ;U       C  c  F`     )#)))X)))r;rZc eZdZdededejjfdZe ddejjdede efdZ d Z dejjddfd Z d ed eddfd ZedefdZdS)BufferShuffledExamplesIterabler buffer_sizerkc0||_||_||_dSr%)rrkrk)rUrrkrks r0rVz'BufferShuffledExamplesIterable.__init__ws&&"r;rrr#c#ZK d|d||DEd{V))NTc34K|]}t|VdSr%rrs r0rzFBufferShuffledExamplesIterable._iter_random_indices..s(]]1A]]]]]]r;rr)r!)rrkr#s r0r$z3BufferShuffledExamplesIterable._iter_random_indices|sK ^]] Q J[ (\(\]]] ] ] ] ] ] ] ] ^r;c#LK|j}t|j}|||}g}|jD]I}t ||krt |}||V|||<4||J|||Ed{VdSr%) rkrrkr$rrIrJrr)rUrkrr mem_bufferxrHs r0rYz'BufferShuffledExamplesIterable.__iter__s& t~&&44S+FF ! % %A:+--)** m### ! 1 !!!$$$$ Jr;c`t|j||j|S)zFShuffle the wrapped examples iterable as well as the shuffling buffer.rkrk)rjrrnrkrms r0rnz3BufferShuffledExamplesIterable.shuffle_data_sourcess5-   1 1) < <$JZfo    r;rorpclt|j|||j|jS)rrs)rjrrtrkrkrss r0rtz1BufferShuffledExamplesIterable.shard_data_sourcess:-   / / ; G G(n    r;c|jjSr%rrXs r0rwz'BufferShuffledExamplesIterable.n_shardsrr;N)r)rarbrcrfr}rzr{r|rVr+r r$rYrnrtr~rwr5r;r0rjrjvs#$9##XZXaXk#### ^^")"5^C^dlmpdq^^^\^" bi.A Ff     C c Ff    )#)))X)))r;rjcdeZdZdedefdZdZdejj ddfdZ e defdZ d S) SkipExamplesIterablernc"||_||_dSr%rrxrUrrxs r0rVzSkipExamplesIterable.__init__&r;c#LKt|j|jdEd{VdSr%rrrxrXs r0rYzSkipExamplesIterable.__iter__s5$*DFD99999999999r;rkr#c|S)zeDoesn't shuffle the wrapped examples iterable since it would skip examples from other shards instead.r5rms r0rnz)SkipExamplesIterable.shuffle_data_sources r;c|jjSr%rrXs r0rwzSkipExamplesIterable.n_shardsrr;N) rarbrcrfr}rVrYrzr{r|rnr~rwr5r;r0rwrws$9c:::bi.AF\)#)))X)))r;rwceZdZdedefdZdZdejj ddfdZ e dZ d ed eddfd Z edefd Zd S)TakeExamplesIterablerrxc"||_||_dSr%rzr{s r0rVzTakeExamplesIterable.__init__r|r;c#JKt|j|jEd{VdSr%r~rXs r0rYzTakeExamplesIterable.__iter__s3$*DF33333333333r;rkr#c|S)zeDoesn't shuffle the wrapped examples iterable since it would take examples from other shards instead.r5rms r0rnz)TakeExamplesIterable.shuffle_data_sourcesrr;cj||z}||z}|g|z}t|D]}||xxdz cc<|Sr)rL)numrxquotient remainderr\rHs r0 split_numberz!TakeExamplesIterable.split_numbersR!8!G ay!!  A 1IIINIIII r;rorpct|j||||j||S)r)rx)rrrtrrxrss r0rtz'TakeExamplesIterable.shard_data_sourcessI#   / / ; G G 44Y?    r;c|jjSr%rrXs r0rwzTakeExamplesIterable.n_shardsrr;N)rarbrcrfr}rVrYrzr{r|rnr+rrtr~rwr5r;r0rrs$9c444bi.AF\\ C c F\    )#)))X)))r;rr8featurestoken_per_repo_idct|}|D] }||vrd||< ||}|||}|SNr)r@encode_exampledecode_example)r8rrrencoded_exampledecoded_examples r0_apply_feature_types_on_examplersf7mmG(( g % %#'GK --g66O--oQb-ccO r;ct|}t|tt|}|D]}||vr dg|z||<||}|||}|Sr)r@rIrJrK encode_batch decode_batch)r!rrrNr encoded_batch decoded_batchs r0_apply_feature_types_on_batchrs KKEU4U ,,-..J55 e # #"&*!4E+ ))%00M))-K\)]]M r;c eZdZdededeeeeedfffdZ dZ de j j ddfd Zd ed eddfd Zedefd ZdS)TypedExamplesIterablerrrNc0||_||_||_dSr%)rrr)rUrrrs r0rVzTypedExamplesIterable.__init__s! '  !2r;c#dK|jD]%\}}|t||j|jfV&dSr)rrrr)rUrr8s r0rYzTypedExamplesIterable.__iter__sb!,  LC6$:P      r;rkr#cjt|j||j|jS)rTrr)rrrnrrrms r0rnz*TypedExamplesIterable.shuffle_data_sources s8$   1 1) < <]"4    r;rorpclt|j|||j|jS)rr)rrrtrrrss r0rtz(TypedExamplesIterable.shard_data_sourcess:$   / / ; G G]"4    r;c|jjSr%rrXs r0rwzTypedExamplesIterable.n_shardsrr;)rarbrcrfrr rrrXrVrYrzr{r|rnr}rtr~rwr5r;r0rrs3*33 U3d?%; ;< 3333 bi.A F]     C c F]    )#)))X)))r;rcfd}|S)Nc;Kt}di|D]G\}}||}tt|D]\}}|d||fVHdS)Nrr5)r format_batchrKrO)rpython_formatterrtabler!rHr8generate_tables_fns r0wrapperz7_generate_examples_from_tables_wrapper..wrappers*,,,,66v66 , ,JC$11%88E'(:5(A(ABB , , 7llqllG+++++ , , ,r;r5)rrs` r0&_generate_examples_from_tables_wrapperrs#,,,,, Nr;cHeZdZUejjed<dZee ed<dS)ShufflingConfigrkN_original_seed) rarbrcrzr{r|__annotations__rr r}r5r;r0rr)s7y""""$(NHSM(((((r;rc$eZdZUeed<eed<dS)DistributedConfigrank world_sizeN)rarbrcr}rr5r;r0rr/s" IIIOOOOOr;rctjr>ddl}|jjj|jvr$|xj|jjjfz c_dSdSdS)zNAdd torch.utils.data.IterableDataset as a parent class if 'torch' is availablerN)rTORCH_AVAILABLEtorch.utils.datautilsdataIterableDataset __bases__)clstorchs r0._maybe_add_torch_iterable_dataset_parent_classr5sc A ;  +3= @ @ MMek.>@ @MMMM AA A @r;c,eZdZdZ dBdedeedeedeedee dee d ee ee ee dfffd Zd Zd ZdCdZdZedefdZdefdZdZdefdZdZdDdede fdZe dEdedeedeeddfdZ dFdeeddfdZ dGd!eed"e d#ee ee!efd$e deede d%ee ee!efdeed&eeddfd'Z" dHd!eed#ee ee!efd$e deeddf d(Z# dIdee$j%j&d)eddfd*Z'd+efd,Z(dJd-Z)dJd.Z*edee!efd/Z+d0ed1e e,e$j-fddfd2Z.d3ed4eddfd5Z/d6e eefddfd7Z0d8e ee!efddfd9Z1d8e ee!efddfd:Z2d1ed;e3ddfd<Z4deddfd=Z5d>ed?eddfd@Z6dAZ7dS)Krz A Dataset backed by an iterable.Nrinfosplit format_type shuffling distributedrcF|r#|jdkr|r|jtd||n t }t j|||||_||_||_ ||_ d|_ |pi|_ t|jdS)NrzThe dataset doesn't have a fixed random seed across nodes to shuffle and split the list of dataset shards by node. Please pass e.g. `seed=42` in `.shuffle()` to make all the nodes use the same seed. )rrr)rr RuntimeErrorcopyrrrV _ex_iterable _format_type _shuffling _distributed_epoch_token_per_repo_idrr)rUrrrrrrrs r0rVzIterableDataset.__init__As  ;1A55)5 H`Hhg  #.tyy{{{KMM!$T????''#' EVE\Z\6t~FFFFFr;c|jSr%)__dict__rXs r0 __getstate__zIterableDataset.__getstate__\s }r;c<||_t|jdSr%)rrr)rUds r0 __setstate__zIterableDataset.__setstate___s 6t~FFFFFr;c`tt||Sr%)rDrvtake)rUrxs r0_headzIterableDataset._headds"!$tyy||"4"4555r;c,|jr|jdkr |jjS|jrat|jjdd|jz }|dkrd|zn|}t j|Std)NrlzThis dataset is not shuffled) rrrkrr!rzr{ default_rngr)rUeffective_seeds r0_effective_generatorz$IterableDataset._effective_generatorgs ? =t{a//?, , _ =%do&?@@II!WUUX\XccN;IA;M;Mg77SaN9((88 8;<< O>Z!Z^_!_!_$-1B1MM M ))r;c #Ktddl}|jj}|rv|j|jkrft d|jd|jd|j|jz dt d|jd|jd|j rd |j j d nd }|j |j|j}|rt|d |jd t!|d|jd||j|jD]/\}}|jrt'||j|jV+|V0t|d |jdt!|d|jddSt|d |jd|jd|jddS)NrzToo many dataloader workers: z (max is dataset.n_shards=z ). Stopping z dataloader workers.zTo parallelize data loading, we give each process some shards (or data sources) to process. Therefore it's unnecessary to have a number of workers greater than dataset.n_shards=zJ. To enable more parallelism, please split the dataset in more files than rHznode# zdataloader worker#z, ': Starting to iterate over /z shards.rz, ': Finished iterating over z9, ': Stopping... Number of dataset shards < num_workers (U'U'U NNa 0Gaacncwaa'3k6JJaaa    NNshsh|ss[f[osss    <@;LT7d/47777RT *HHYdYpqq   LLUU+.UU`cdr`s`sUUwBwKUUU   !, > >{~{Of g g " " W="9$BY"MMMM LLTT+.TT_bcq_r_rTTvAvJTTT      LLnn+.nn|G|PnnS^Sjnnn     r;c|jr|jjdkrdSdtjvr1ddl}|jj}| |jdkrdSdS)NrFrT) rrsysmodulesrrrrr)rUrrs r0rz IterableDataset._is_main_processsm   !2!7!!;!;5 ck ! ! # # # #+*::<A+=+=utr;c t|jr-|j|}n|j}|jr|jj}|jj}|j|zdkrc|r8|j|z}|dkrdnd}t d|d|d|d| ||}nl|rFt d |d t d |d |jd |t|||}|S)Nrrsrz Assigning z shardz (or data sourcez) of the dataset to each node.zAssigning 1 out of zS examples of the dataset to each node. The others are skipped during the iteration.zIt is more optimized to distribute the dataset shards (or data sources) across nodes. You can do that by using a dataset with number of shards that is a factor of world_size=z. The current dataset has z which is not a factor of r)rrrnrrrrrwrrrrtrr)rUrrrn_shards_per_nodeplurals r0"_prepare_ex_iterable_for_iterationz2IterableDataset._prepare_ex_iterable_for_iterations ? ,+@@AZAZA\A\]]KK+K   ^$)D*5J#j0A55((**(3(< (J%$5$9$9SSrFNN|%6||f||V\|||*<3Gppcmpp 3;ZX\]]] r;c#rK|}dtjvr`ddl}|jj}t||jjjr|| |Ed{VdS|D]/\}}|j rt||j |j V+|V0dS)Nrrr) rrrrrrr isinstancerrrrr)rUrrrrr8s r0rYzIterableDataset.__iter__s==?? ck ! ! # # # #+*::<U   r;Fr2r3c#@Kt|}|D]w\}}|gdt||dz Dz}|rt||krdSt |}|jrt ||j|jVs|VxdS)aIterate through the batches of size `batch_size`. Args: batch_size (:obj:`int`): size of each batch to yield. drop_last_batch (:obj:`bool`, default `False`): Whether a last batch smaller than the batch_size should be dropped cg|]\}}|Sr5r5r9s r0r?z(IterableDataset.iter..s#a#a#a WG#a#a#ar;rNr)rKrrrIrDrrr)rUr2r3rLrr8r2r!s r0rKzIterableDataset.iters??AABB$  LCy#a#axQ[^_Q_@`@`#a#a#aaH 3x==:#=#=&x00E} 4E4=\`\sttttttt   r;rkr gen_kwargscPddlm}||||dS)aCreate an Iterable Dataset from a generator. Args: generator (`Callable`): A generator function that `yields` examples. features (`Features`, *optional*): Dataset features. gen_kwargs(`dict`, *optional*): Keyword arguments to be passed to the `generator` callable. You can define a sharded iterable dataset by passing the list of shards in `gen_kwargs`. This can be used to improve shuffling and when iterating over the dataset with multiple workers. Returns: `IterableDataset` Example: ```py >>> def gen(): ... yield {"text": "Good", "label": 0} ... yield {"text": "Bad", "label": 1} ... >>> ds = IterableDataset.from_generator(gen) ``` ```py >>> def gen(shards): ... for shard in shards: ... with open(shard) as f: ... for line in f: ... yield {"line": line} ... >>> shards = [f"data{i}.txt" for i in range(32)] >>> ds = IterableDataset.from_generator(gen, gen_kwargs={"shards": shards}) >>> ds = ds.shuffle(seed=42, buffer_size=10_000) # shuffles the shards order + uses a shuffle buffer >>> from torch.utils.data import DataLoader >>> dataloader = DataLoader(ds.with_format("torch"), num_workers=4) # give each worker a subset of 32/4=8 shards ``` r)GeneratorDatasetInputStreamT)rkrr streaming) io.generatorrread)rkrrrs r0from_generatorzIterableDataset.from_generatorsHZ >=====**!    $&&  r;r*c t|}t|j|j|j|t j|jt j|j|j S)aG Return a dataset with the specified format. This method only supports the "torch" format for now. Args: type (`str`, optional, default None): if set to "torch", the returned dataset will be a subclass of torch.utils.data.IterableDataset to be used in a DataLoader rrrrrrr) rrr_infor_splitrrrr)rUr*s r0 with_formatzIterableDataset.with_format+sl*$// )""+mDO44 d&788"5    r;rr.r/r0r1r4r5c t|tr|g}t|tr|g}|d}| i} t|jj&t |j|jj|jn|j||||||||  } |j } || _t| | |j |j tj |jtj |j|jS)a Apply a function to all the examples in the iterable dataset (individually or in batches) and update them. If your function returns a column that already exists, then it overwrites it. The function is applied on-the-fly on the examples when iterating over the dataset. You can specify whether the function should be batched or not with the `batched` parameter: - If batched is `False`, then the function takes 1 example in and should return 1 example. An example is a dictionary, e.g. `{"text": "Hello there !"}`. - If batched is `True` and `batch_size` is 1, then the function takes a batch of 1 example as input and can return a batch with 1 or more examples. A batch is a dictionary, e.g. a batch of 1 example is {"text": ["Hello there !"]}. - If batched is `True` and `batch_size` is `n` > 1, then the function takes a batch of `n` examples as input and can return a batch with `n` examples, or with an arbitrary number of examples. Note that the last batch may have less than `n` examples. A batch is a dictionary, e.g. a batch of `n` examples is `{"text": ["Hello there !"] * n}`. Args: function (`Callable`, *optional*, defaults to `None`): Function applied on-the-fly on the examples when you iterate on the dataset. It must have one of the following signatures: - `function(example: Dict[str, Any]) -> Dict[str, Any]` if `batched=False` and `with_indices=False` - `function(example: Dict[str, Any], idx: int) -> Dict[str, Any]` if `batched=False` and `with_indices=True` - `function(batch: Dict[str, List]) -> Dict[str, List]` if `batched=True` and `with_indices=False` - `function(batch: Dict[str, List], indices: List[int]) -> Dict[str, List]` if `batched=True` and `with_indices=True` For advanced usage, the function can also return a `pyarrow.Table`. Moreover if your function returns nothing (`None`), then `map` will run your function and return the dataset unchanged. If no function is provided, default to identity function: `lambda x: x`. with_indices (`bool`, defaults to `False`): Provide example indices to `function`. Note that in this case the signature of `function` should be `def function(example, idx[, rank]): ...`. input_columns (`Optional[Union[str, List[str]]]`, defaults to `None`): The columns to be passed into `function` as positional arguments. If `None`, a dict mapping to all formatted columns is passed as one argument. batched (`bool`, defaults to `False`): Provide batch of examples to `function`. batch_size (`int`, *optional*, defaults to `1000`): Number of examples per batch provided to `function` if `batched=True`. `batch_size <= 0` or `batch_size == None` then provide the full dataset as a single batch to `function`. drop_last_batch (`bool`, defaults to `False`): Whether a last batch smaller than the batch_size should be dropped instead of being processed by the function. remove_columns (`[List[str]]`, *optional*, defaults to `None`): Remove a selection of columns while doing the mapping. Columns will be removed before updating the examples with the output of `function`, i.e. if `function` is adding columns with names in `remove_columns`, these columns will be kept. features (`[Features]`, *optional*, defaults to `None`): Feature types of the resulting dataset. fn_kwargs (`Dict`, *optional*, default `None`): Keyword arguments to be passed to `function`. Example: ```py >>> from datasets import load_dataset >>> ds = load_dataset("rotten_tomatoes", split="train", streaming=True) >>> def add_prefix(example): ... example["text"] = "Review: " + example["text"] ... return example >>> ds = ds.map(add_prefix) >>> list(ds.take(3)) [{'label': 1, 'text': 'Review: the rock is destined to be the 21st century's new " conan " and that he's going to make a splash even greater than arnold schwarzenegger , jean-claud van damme or steven segal .'}, {'label': 1, 'text': 'Review: the gorgeously elaborate continuation of " the lord of the rings " trilogy is so huge that a column of words cannot adequately describe co-writer/director peter jackson's expanded vision of j . r . r . tolkien's middle-earth .'}, {'label': 1, 'text': 'Review: effective but too-tepid biopic'}] ``` Nc|Sr%r5)rqs r0z%IterableDataset.map..sr;r)r.r/r0r1r2r3r4r5r)rrr-rrrrrrrrrrrrr) rUr.r/r0r1r2r3r4rr5rrs r0mapzIterableDataset.mapGs^ mS ) ) ,*OM nc * * .,-N  "{H  I,z". "$"3TZ5H\`\s t t t t"%'!+)    y~~  #+)mDO44 d&788"5    r;c t|tr|g}tj|j}d|_t |jj&t|j|jj|j n|j|||||}t|||j |j tj|j tj|j|j S)aApply a filter function to all the elements so that the dataset only includes examples according to the filter function. The filtering is done on-the-fly when iterating over the dataset. Args: function (`Callable`): Callable with one of the following signatures: - `function(example: Dict[str, Any]) -> bool` if `with_indices=False, batched=False` - `function(example: Dict[str, Any], indices: int) -> bool` if `with_indices=True, batched=False` - `function(example: Dict[str, List]) -> List[bool]` if `with_indices=False, batched=True` - `function(example: Dict[str, List], indices: List[int]) -> List[bool]` if `with_indices=True, batched=True` If no function is provided, defaults to an always True function: `lambda x: True`. with_indices (`bool`, defaults to `False`): Provide example indices to `function`. Note that in this case the signature of `function` should be `def function(example, idx): ...`. input_columns (`str` or `List[str]`, *optional*): The columns to be passed into `function` as positional arguments. If `None`, a dict mapping to all formatted columns is passed as one argument. batched (`bool`, defaults to `False`): Provide batch of examples to `function`. batch_size (`int`, *optional*, default `1000`): Number of examples per batch provided to `function` if `batched=True`. Example: ```py >>> from datasets import load_dataset >>> ds = load_dataset("rotten_tomatoes", split="train", streaming=True) >>> ds = ds.filter(lambda x: x["label"] == 0) >>> list(ds.take(3)) [{'label': 0, 'movie_review': 'simplistic , silly and tedious .'}, {'label': 0, 'movie_review': "it's so laddish and juvenile , only teenage boys could possibly find it funny ."}, {'label': 0, 'movie_review': 'exploitative and largely devoid of the depth or sophistication that would make watching such a graphic treatment of the crimes bearable .'}] ``` Nrrfr)rrrrrrrZrrrrrrrr)rUr.r/r0r1r2rrs r0filterzIterableDataset.filtersZ mS ) ) ,*OM}TZ(( /z". "$"3TZ5H\`\s t t t t"%'!    #+)mDO44 d&788"5    r;rkc || tj|}nt|}t ||}t t |j||||j |j |j |tj|j |jS)aC Randomly shuffles the elements of this dataset. This dataset fills a buffer with `buffer_size` elements, then randomly samples elements from this buffer, replacing the selected elements with new elements. For perfect shuffling, a buffer size greater than or equal to the full size of the dataset is required. For instance, if your dataset contains 10,000 elements but `buffer_size` is set to 1000, then `shuffle` will initially select a random element from only the first 1000 elements in the buffer. Once an element is selected, its space in the buffer is replaced by the next (i.e. 1,001-st) element, maintaining the 1000 element buffer. If the dataset is made of several shards, it also does shuffle the order of the shards. However if the order has been fixed by using [`~datasets.IterableDataset.skip`] or [`~datasets.IterableDataset.take`] then the order of the shards is kept unchanged. Args: seed (`int`, *optional*, defaults to `None`): Random seed that will be used to shuffle the dataset. It is used to sample from the shuffle buffe and also to shuffle the data shards. generator (`numpy.random.Generator`, *optional*): Numpy random Generator to use to compute the permutation of the dataset rows. If `generator=None` (default), uses `np.random.default_rng` (the default BitGenerator (PCG64) of NumPy). buffer_size (`int`, defaults to `1000`): Size of the buffer. Example: ```py >>> from datasets import load_dataset >>> ds = load_dataset("rotten_tomatoes", split="train", streaming=True) >>> list(ds.take(3)) [{'label': 1, 'text': 'the rock is destined to be the 21st century's new " conan " and that he's going to make a splash even greater than arnold schwarzenegger , jean-claud van damme or steven segal .'}, {'label': 1, 'text': 'the gorgeously elaborate continuation of " the lord of the rings " trilogy is so huge that a column of words cannot adequately describe co-writer/director peter jackson's expanded vision of j . r . r . tolkien's middle-earth .'}, {'label': 1, 'text': 'effective but too-tepid biopic'}] >>> shuffled_ds = ds.shuffle(seed=42) >>> list(shuffled_ds.take(3)) [{'label': 1, 'text': "a sports movie with action that's exciting on the field and a story you care about off it ."}, {'label': 1, 'text': 'at its best , the good girl is a refreshingly adult take on adultery . . .'}, {'label': 1, 'text': "sam jones became a very lucky filmmaker the day wilco got dropped from their record label , proving that one man's ruin may be another's fortune ."}] ``` N)rkrrsr)rzr{rrrrrjrrnrrrrrr)rUrdrkrkrs r0rzIterableDataset.shufflesd   --d33II ++I#iMMM 6!{i""9--""+) d&788"5    r;epochc||_dSr%)r)rUr s r0 set_epochzIterableDataset.set_epochCs  r;c t|j|}t||j|j|jt j|jt j|j |j S)a Create a new [`IterableDataset`] that skips the first `n` elements. Args: n (`int`): Number of elements to skip. Example: ```py >>> from datasets import load_dataset >>> ds = load_dataset("rotten_tomatoes", split="train", streaming=True) >>> list(ds.take(3)) [{'label': 1, 'text': 'the rock is destined to be the 21st century's new " conan " and that he's going to make a splash even greater than arnold schwarzenegger , jean-claud van damme or steven segal .'}, {'label': 1, 'text': 'the gorgeously elaborate continuation of " the lord of the rings " trilogy is so huge that a column of words cannot adequately describe co-writer/director peter jackson's expanded vision of j . r . r . tolkien's middle-earth .'}, {'label': 1, 'text': 'effective but too-tepid biopic'}] >>> ds = ds.skip(1) >>> list(ds.take(3)) [{'label': 1, 'text': 'the gorgeously elaborate continuation of " the lord of the rings " trilogy is so huge that a column of words cannot adequately describe co-writer/director peter jackson's expanded vision of j . r . r . tolkien's middle-earth .'}, {'label': 1, 'text': 'effective but too-tepid biopic'}, {'label': 1, 'text': 'if you sometimes like to go to the movies to have fun , wasabi is a good place to start .'}] ``` r) rwrrrrrrrrrrrUrxrs r0skipzIterableDataset.skipFso8+4+>> from datasets import load_dataset >>> ds = load_dataset("rotten_tomatoes", split="train", streaming=True) >>> small_ds = ds.take(2) >>> list(small_ds) [{'label': 1, 'text': 'the rock is destined to be the 21st century's new " conan " and that he's going to make a splash even greater than arnold schwarzenegger , jean-claud van damme or steven segal .'}, {'label': 1, 'text': 'the gorgeously elaborate continuation of " the lord of the rings " trilogy is so huge that a column of words cannot adequately describe co-writer/director peter jackson's expanded vision of j . r . r . tolkien's middle-earth .'}] ``` r) rrrrrrrrrrrrs r0rzIterableDataset.takemso*+4+>> from datasets import load_dataset >>> ds = load_dataset("rotten_tomatoes", split="validation", streaming=True) >>> ds.column_names ['text', 'label'] ``` N)rrrvr rXs r0rzIterableDataset.column_namess348:3F3RtDJ',,..///X\\r;namecolumnc@fd}||dS)zAdd column to Dataset. Args: name (str): Column name. column (list or np.array): Column data to be added. Returns: `IterableDataset` cL|vrtddd|iS)NzError when adding : column  is already in the dataset.r)r8rrrs r0 add_column_fnz1IterableDataset.add_column..add_column_fns=w !fd!f!fT!f!f!fggg&+& &r;T)r/)r)rUrrrs `` r0 add_columnzIterableDataset.add_columns: ' ' ' ' ' ' xx Dx999r;original_column_namenew_column_namecfd}|jjr|jjnd}||g}|8t fd|D|j_|S)ax Rename a column in the dataset, and move the features associated to the original column under the new column name. Args: original_column_name (`str`): Name of the column to rename. new_column_name (`str`): New name for the column. Returns: `IterableDataset`: A copy of the dataset with a renamed column. Example: ```py >>> from datasets import load_dataset >>> ds = load_dataset("rotten_tomatoes", split="train", streaming=True) >>> next(iter(ds)) {'label': 1, 'text': 'the rock is destined to be the 21st century's new " conan " and that he's going to make a splash even greater than arnold schwarzenegger , jean-claud van damme or steven segal .'} >>> ds = ds.rename_column("text", "movie_review") >>> next(iter(ds)) {'label': 1, 'movie_review': 'the rock is destined to be the 21st century's new " conan " and that he's going to make a splash even greater than arnold schwarzenegger , jean-claud van damme or steven segal .'} ``` c |vrtdddd|vrtdddd|iS)NError when renaming  to rz is not in the dataset.rr)r8rrs r0rename_column_fnz7IterableDataset.rename_column..rename_column_fns#722 M+?MM_MM_sMMM')) L+?LL_LL_nLLL$W-A%BC Cr;Nr4c,i|]\}}|krn||Sr5r5)r7r9featurerrs r0r:z1IterableDataset.rename_column..s>$W(+.B'B'BOOWr;)rrrrrrM)rUrrr#original_features ds_iterables `` r0 rename_columnzIterableDataset.rename_columns: D D D D D D;?*:MWDJ/44666SWhh/AU@VhWW  ()1(9(?(?(A(A**K  & r;column_mappingc$fd}|jjr|jjnd}||t }|7t fd|D|j_|S)aa Rename several columns in the dataset, and move the features associated to the original columns under the new column names. Args: column_mapping (`Dict[str, str]`): A mapping of columns to rename to their new names Returns: `IterableDataset`: A copy of the dataset with renamed columns c tfdDrbtdtdtdt t z dtfdDrttdtdtdt t z dfdDS) Nc3 K|]}|vV dSr%r5r7r9r8s r0rzLIterableDataset.rename_columns..rename_columns_fn..s(@@#3g%@@@@@@r;r!r"z : columns z are not in the dataset.c3 K|]}|vV dSr%r5r.s r0rzLIterableDataset.rename_columns..rename_columns_fn..s'EEc3'>EEEEEEr;z are already in the dataset.c(i|]\}}||Sr5r5)r7rrr8s r0r:zMIterableDataset.rename_columns..rename_columns_fn..s69(/ )=!>r;)rrrvrsetrM)r8r*s`r0rename_columns_fnz9IterableDataset.rename_columns..rename_columns_fns@@@@@@@@@  k4+?+?kkT.J_J_JaJaEbEbkknqsAoBoBEHIPEQEQoQkkkEEEE^-B-B-D-DEEEEE  x4+?+?xxT.J_J_JaJaEbEbxxnqrynznz~ABPBWBWBYBY~Z~ZoZxxx=K=Q=Q=S=S r;Nr$cXi|]&\}}|vr|n||'Sr5)r )r7r9r&r*s r0r:z2IterableDataset.rename_columns..sP$W,/.2E2E2G2G+G+GN3''SRYr;)rrrrrvrrM)rUr*r2r'r(s ` r0rename_columnszIterableDataset.rename_columnss     ;?*:MWDJ/44666SWhh0nAUAUhVV  ()1(9(?(?(A(A**K  & r;rc |jjr|jjnd}||}|I||j_|D]\}}||vr |jj|=|S)aD Remove one or several column(s) in the dataset and the features associated to them. The removal is done on-the-fly on the examples when iterating over the dataset. Args: column_names (`Union[str, List[str]]`): Name of the column(s) to remove. Returns: `IterableDataset`: A copy of the dataset object without the columns to remove. Example: ```py >>> from datasets import load_dataset >>> ds = load_dataset("rotten_tomatoes", split="train", streaming=True) >>> next(iter(ds)) {'text': 'the rock is destined to be the 21st century's new " conan " and that he's going to make a splash even greater than arnold schwarzenegger , jean-claud van damme or steven segal .', 'label': 1} >>> ds = ds.remove_columns("label") >>> next(iter(ds)) {'text': 'the rock is destined to be the 21st century's new " conan " and that he's going to make a splash even greater than arnold schwarzenegger , jean-claud van damme or steven segal .'} ``` Nr$)rrrrrM)rUrr'r(r9rs r0r4zIterableDataset.remove_columnss2;?*:MWDJ/44666SWhhlh;;  ():)?)?)A)AK  &+1133 8 8Q,&&#)237r;c t|tr|g}|jrtj|j|jjr|D]O}||jjvr?t d|dt|jjdPtfd|D_t|j |}t||j |j|j|j|jS)a<Select one or several column(s) in the dataset and the features associated to them. The selection is done on-the-fly on the examples when iterating over the dataset. Args: column_names (`Union[str, List[str]]`): Name of the column(s) to select. Returns: `IterableDataset`: A copy of the dataset object with selected columns. Example: ```py >>> from datasets import load_dataset >>> ds = load_dataset("rotten_tomatoes", split="train", streaming=True) >>> next(iter(ds)) {'text': 'the rock is destined to be the 21st century's new " conan " and that he's going to make a splash even greater than arnold schwarzenegger , jean-claud van damme or steven segal .', 'label': 1} >>> ds = ds.select_columns("text") >>> next(iter(ds)) {'text': 'the rock is destined to be the 21st century's new " conan " and that he's going to make a splash even greater than arnold schwarzenegger , jean-claud van damme or steven segal .'} ``` Nz Column name z- not in the dataset. Columns in the dataset: rHc,i|]}|j|Sr5r)r7rrs r0r:z2IterableDataset.select_columns..Ns")T)T)T!!T]1-=)T)T)Tr;r)rrrrrrrrvr rrrrrrrrr)rUrrrrs @r0select_columnszIterableDataset.select_columns(s.2 lC ( ( *(>L : V=,,Dz".#/K"$**===(C;CC#DJ$7$<$<$>$>??CCC> !))T)T)T)T|)T)T)T U U +D,=|LL #+)o)"5    r;r&c R|j}||j|< |n#t$r d|_YnwxYwt |j||j|jtj |j tj |j |j S)aCast column to feature for decoding. Args: column (`str`): Column name. feature (`Feature`): Target feature. Returns: `IterableDataset` Example: ```py >>> from datasets import load_dataset, Audio >>> ds = load_dataset("PolyAI/minds14", name="en-US", split="train", streaming=True) >>> ds.features {'audio': Audio(sampling_rate=8000, mono=True, decode=True, id=None), 'english_transcription': Value(dtype='string', id=None), 'intent_class': ClassLabel(num_classes=14, names=['abroad', 'address', 'app_error', 'atm_limit', 'balance', 'business_loan', 'card_issues', 'cash_deposit', 'direct_debit', 'freeze', 'high_value_payment', 'joint_account', 'latest_transactions', 'pay_bill'], id=None), 'lang_id': ClassLabel(num_classes=14, names=['cs-CZ', 'de-DE', 'en-AU', 'en-GB', 'en-US', 'es-ES', 'fr-FR', 'it-IT', 'ko-KR', 'nl-NL', 'pl-PL', 'pt-PT', 'ru-RU', 'zh-CN'], id=None), 'path': Value(dtype='string', id=None), 'transcription': Value(dtype='string', id=None)} >>> ds = ds.cast_column("audio", Audio(sampling_rate=16000)) >>> ds.features {'audio': Audio(sampling_rate=16000, mono=True, decode=True, id=None), 'english_transcription': Value(dtype='string', id=None), 'intent_class': ClassLabel(num_classes=14, names=['abroad', 'address', 'app_error', 'atm_limit', 'balance', 'business_loan', 'card_issues', 'cash_deposit', 'direct_debit', 'freeze', 'high_value_payment', 'joint_account', 'latest_transactions', 'pay_bill'], id=None), 'lang_id': ClassLabel(num_classes=14, names=['cs-CZ', 'de-DE', 'en-AU', 'en-GB', 'en-US', 'es-ES', 'fr-FR', 'it-IT', 'ko-KR', 'nl-NL', 'pl-PL', 'pt-PT', 'ru-RU', 'zh-CN'], id=None), 'path': Value(dtype='string', id=None), 'transcription': Value(dtype='string', id=None)} ``` Nr rrrrtask_templatesrrrrrrrr)rUrr&rs r0 cast_columnzIterableDataset.cast_column[sDz   ' f ' IIKKKK ' ' '"&D    ')+)mDO44 d&788"5    s:A Ac L|j}||_ |n#t$r d|_YnwxYwt |j||j|jtj |j tj |j |j S)a Cast the dataset to a new set of features. Args: features ([`Features`]): New features to cast the dataset to. The name of the fields in the features must match the current column names. The type of the data must also be convertible from one type to the other. For non-trivial conversion, e.g. `string` <-> `ClassLabel` you should use [`~Dataset.map`] to update the Dataset. Returns: `IterableDataset`: A copy of the dataset with casted features. Example: ```py >>> from datasets import load_dataset >>> ds = load_dataset("rotten_tomatoes", split="train", streaming=True) >>> ds.features {'label': ClassLabel(num_classes=2, names=['neg', 'pos'], id=None), 'text': Value(dtype='string', id=None)} >>> new_features = ds.features.copy() >>> new_features["label"] = ClassLabel(names=["bad", "good"]) >>> new_features["text"] = Value("large_string") >>> ds = ds.cast(new_features) >>> ds.features {'label': ClassLabel(num_classes=2, names=['bad', 'good'], id=None), 'text': Value(dtype='large_string', id=None)} ``` Nrr;rUrrs r0castzIterableDataset.castsDz     ' IIKKKK ' ' '"&D    ')+)mDO44 d&788"5    s7A  A rrc t|j||}t||j|j|jt j|jt j|j |j S)Nrr) rrrrrrrrrrr)rUrrrs r0_stepzIterableDataset._stepss*4+<4PVWWW #""+)mDO44 d&788"5    r;c |j|St|jtr |jj}n!t |}|j}||_t|j||j |j tj |j tj |j |jS)Nr)rrrrr1rrrrrrrrrrr?s r0_resolve_featuresz!IterableDataset._resolve_featuress = $K )+@ A A @(1HH1$**,,??Hy~~  )+)mDO44 d&788"5    r;)NNNNNN)r)F)NNr%) NFNFrFNNN)NFNFr)NNr)r#r)8rarbrcrdrfr rrrrrr rrXrVrrrrr~r}rwrrrrYrKr+r rr@rrr rr rzr{r|rrrrrrvrGrr)r4r4r9rr=r@rBrDr5r;r0rr>s** '+&*%)/337IMGG*G{#G # G c] G O, G/0G$DeCtO.D)D$EFGGGG6GGG 6666 = = =*#***X* %)>%%%%N   4I>,sT,(,%)3338$3TN3  333\3n#  sm      <(,"9=$( %:>'+$(n n 8$n n  c49n 56 n  n SM n n !sDI~!67n 8$n D>n  n n n n d(,9=$( G G 8$G  c49n 56 G  G SM G  G G G G T^bA A $,RY-@$AA WZA A A A A Fs% % % % N    @ ]htCy1 ] ] ]X ]:s:E$.,A:FW::::$1#11Pa1111f#T#s(^#@Q####J 5d3i+@ EV    D1 5d3i+@1 EV1 1 1 1 f1 #1 1 @Q1 1 1 1 f1 1  1 1 1 1 f  #  s  /@          r;rdsetsrraxiscd|D}|dkrtd|Dntd|Dtdtd|DD}d|D}|dkrt |}nt |}|t jd |D}n|}||_ d |D}t|||| S) a Converts a list of `IterableDataset` with the same schema into a single `IterableDataset`. Missing data are filled with None values. Args: dsets (`List[datasets.IterableDataset]`): List of Datasets to concatenate. info (`DatasetInfo`, optional): Dataset information, like description, citation, etc. split (`NamedSplit`, optional): Name of the dataset split. axis (``{0, 1}``, default ``0``, meaning over rows): Axis to concatenate over, where ``0`` means over rows (vertically) and ``1`` means over columns (horizontally). *New in version 1.6.0* Example: ```py >>> ds3 = _concatenate_iterable_datasets([ds1, ds2]) ``` c6g|]}|Sr5rDr7rs r0r?z2_concatenate_iterable_datasets..s$ 2 2 2qQ " " 2 2 2r;rcg|] }|j Sr5r8r7dsets r0r?z2_concatenate_iterable_datasets..s*K*K*KT4=*K*K*Kr;c&g|]}|jD]}|Sr5r8)r7rMcol_names r0r?z2_concatenate_iterable_datasets..s'VVV$ VVHXVVVVr;cHi|]}|D]\}}|| Sr5rMr7rkvs r0r:z2_concatenate_iterable_datasets..s;nnn(]e]k]k]m]mnnUYUVXYAnnnnr;cg|] }|j Sr5r8rLs r0r?z2_concatenate_iterable_datasets..s.O.O.Ot}.O.O.Or;cg|] }|j Sr5rrJs r0r?z2_concatenate_iterable_datasets.. s222qAN222r;Ncg|] }|j Sr5rrJs r0r?z2_concatenate_iterable_datasets..s&=&=&=!qv&=&=&=r;cRi|]$}|jD]\}}|| %Sr5rrMr7datasetrepo_idtokens r0r:z2_concatenate_iterable_datasets..s;vvvGSZSmSsSsSuSuvv%%vvvvr;rrrr) rrrrrrr from_mergerrr)rErrrFrrrrs r0_concatenate_iterable_datasetsrbs:8 3 2E 2 2 2E qyy)*K*KU*K*K*KLLLLVV%VVVWWWnno.O.O.O.O.OPPnnnH32E222L qyyHVV J<XX  |%&=&=u&=&=&=>>yy{{DMvvuvvv {U^o p p ppr;rdatasetsrrdrcd|D}td|Dtdtd|DD}d|D}|t||}n2tj|} t|| ||}|tj d |D}n| }||_ d |D} t|||| S) a Interleave several iterable datasets (sources) into a single iterable dataset. The new iterable dataset alternates between the sources to yield examples. If `probabilities = None` (default) the iterable dataset will cycles through the sources in order for each next example in the iteration. If `probabilities` is not `None, the iterable dataset will sample a random source according to the provided probabilities for each next examples in the iteration. Args: datasets (`List[IterableDataset]`): list of datasets to interleave probabilities (`List[float]`, optional, default None): If specified, the new iterable dataset samples examples from one source at a time according to these probabilities. seed (`int`, optional, default None): The random seed used to choose a source for each example. stopping_strategy (Optional `str`, defaults to `first_exhausted`): Two strategies are proposed right now. By default, `first_exhausted` is an undersampling strategy, i.e the dataset construction is stopped as soon as one dataset has ran out of samples. If the strategy is `all_exhausted`, we use an oversampling strategy, i.e the dataset construction is stopped as soon as every samples of every dataset has been added at least once. Note that if the strategy is `all_exhausted`, the interleaved dataset size can get enormous: - with no probabilities, the resulting dataset will have max_length_datasets*nb_dataset samples. - with given probabilities, the resulting dataset will have more samples if some datasets have really low probability of visiting. Output: `datasets.IterableDataset` c6g|]}|Sr5rIrJs r0r?z1_interleave_iterable_datasets..=s$888!##%%888r;cg|] }|j Sr5r8rLs r0r?z1_interleave_iterable_datasets..@s&J&J&Jt}&J&J&Jr;cHi|]}|D]\}}|| Sr5rQrRs r0r:z1_interleave_iterable_datasets..Es;qqq(`h`n`n`p`pqqX\XY[\Aqqqqr;cg|] }|j Sr5r8rLs r0r?z1_interleave_iterable_datasets..Es.R.R.Rt}.R.R.Rr;cg|] }|j Sr5rWrJs r0r?z1_interleave_iterable_datasets..Hs555qAN555r;Nr)rkrrcg|] }|j Sr5rYrJs r0r?z1_interleave_iterable_datasets..Us&@&@&@!qv&@&@&@r;cRi|]$}|jD]\}}|| %Sr5r[r\s r0r:z1_interleave_iterable_datasets..ZsN"WE_EeEeEgEg3A7Er;r`) rrrrrzr{rrrrarrr) rcrrdrrrrrrrkrs r0_interleave_iterable_datasetsrlsJ@98x888H&&J&J&J&J&JKKKqqo.R.R.R.R.RSSqqqH65H555L9,Zklll I))$// A I]^o    |%&@&@x&@&@&@AAyy{{DM&. {U^o p p ppr;r]rrc $|jr!||jjz}||jjz|z}t||}t |j|j|j|j tj |j ||j S)a Split an iterable dataset for the node at rank `rank` in a pool of nodes of size `world_size`. If the dataset has a number of shards that is a factor of `world_size` (i.e. if `dataset.n_shards % world_size == 0`), then the shards are evenly assigned across the nodes, which is the most optimized. Otherwise, each node keeps 1 example out of `world_size`, skipping the other examples. Args: dataset ([`IterableDataset`]): The iterable dataset to split by node. rank (`int`): Rank of the current node. world_size (`int`): Total number of nodes. Returns: [`IterableDataset`]: The iterable dataset to be used on the node at rank `rank`. )rrr) rrrrrrrrrrrrr)r]rrrs r0_split_by_node_iterable_datasetrnas&='"6"AA G055<#*EEEK ( ]   ! !n(- 233!4   r;r%)NNr)NNNNr)Xrrr collectionsrr dataclassesrrrtypingrr r r r r rnumpyrzpyarrowr&rr arrow_datasetrrrfeatures.featuresrrr filesystemsr formattingrrrrsplitsrrr utils.loggingrutils.shardingrrrr rarrrvr1rDrOrQrfrrrrrrrrrr-rZrjrwrr@rXrrrrrrrrr}rbr*rlrnr5r;r0r{s  !!!!!!########GGGGGGGGGGGGGGGGGG++++++^^^^^^^^^^++++++CCCCCCCC%%%%%%vvvvvvvvvvvv H  77d39o7XhEW7ck7777#d38n!5#$sDy/####>d39o>$tCH~2F>>>>h:RRRRRRRR4<<<<<,<<<.     *:   ()))))1)))()))))0)))<: : : : : *?: : : z( ( ( ( ( 9N( ( ( V d3i    : : : : : ;P: : : z. . . . . 2U. . . bw)w)w)w)w)2w)w)w)tN)N)N)N)N)4N)N)N)b,),),),),)%:,),),)^)))))0)))")))))0)))B   % :>sE#tUY/DZ?Z:[         # 8