'''simple docstring'''




# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import ast
import difflib
from collections.abc import Mapping
from typing import Any, Callable, Dict




class       __UpperCAmelCase							(							_UpperCamelCase	):





       '''simple docstring'''



       pass





def 	_UpperCAmelCase       (				_UpperCamelCase       :      str,  _UpperCamelCase       :      Dict[str, Callable],  _UpperCamelCase       :      Any=None,  _UpperCamelCase       :      List[str]=False  )      -> str:
    try:
        A_      =  ast.parse(_UpperCamelCase  )
    except SyntaxError as e:
        print('''The code generated by the agent is not valid.\n''',  _UpperCamelCase  )
        return
    if state is None:
        A_      =  {}
    A_      =  None
    for idx, node in enumerate(expression.body  ):
        try:
            A_      =  evaluate_ast(_UpperCamelCase,  _UpperCamelCase,  _UpperCamelCase  )
        except InterpretorError as e:
            A_      =  F'''Evaluation of the code stopped at line {idx} before the end because of the following error'''
            if chat_mode:
                msg += (
                    F'''. Copy paste the following error message and send it back to the agent:\nI get an error: \'{e}\''''
                )
            else:
                msg += F''':\n{e}'''
            print(_UpperCamelCase  )
            break
        if line_result is not None:
            A_      =  line_result

    return result





def 	_UpperCAmelCase       (				_UpperCamelCase       :      ast.AST,  _UpperCamelCase       :      Dict[str, Any],  _UpperCamelCase       :      Dict[str, Callable]  )      -> List[Any]:
    if isinstance(_UpperCamelCase,  ast.Assign  ):
        # Assignement -> we evaluate the assignement which should update the state
        # We return the variable assigned as it may be used to determine the final result.
        return evaluate_assign(_UpperCamelCase,  _UpperCamelCase,  _UpperCamelCase  )
    elif isinstance(_UpperCamelCase,  ast.Call  ):
        # Function call -> we return the value of the function call
        return evaluate_call(_UpperCamelCase,  _UpperCamelCase,  _UpperCamelCase  )
    elif isinstance(_UpperCamelCase,  ast.Constant  ):
        # Constant -> just return the value
        return expression.value
    elif isinstance(_UpperCamelCase,  ast.Dict  ):
        # Dict -> evaluate all keys and values
        A_      =  [evaluate_ast(_UpperCamelCase,  _UpperCamelCase,  _UpperCamelCase  ) for k in expression.keys]
        A_      =  [evaluate_ast(_UpperCamelCase,  _UpperCamelCase,  _UpperCamelCase  ) for v in expression.values]
        return dict(zip(_UpperCamelCase,  _UpperCamelCase  )  )
    elif isinstance(_UpperCamelCase,  ast.Expr  ):
        # Expression -> evaluate the content
        return evaluate_ast(expression.value,  _UpperCamelCase,  _UpperCamelCase  )
    elif isinstance(_UpperCamelCase,  ast.For  ):
        # For loop -> execute the loop
        return evaluate_for(_UpperCamelCase,  _UpperCamelCase,  _UpperCamelCase  )
    elif isinstance(_UpperCamelCase,  ast.FormattedValue  ):
        # Formatted value (part of f-string) -> evaluate the content and return
        return evaluate_ast(expression.value,  _UpperCamelCase,  _UpperCamelCase  )
    elif isinstance(_UpperCamelCase,  ast.If  ):
        # If -> execute the right branch
        return evaluate_if(_UpperCamelCase,  _UpperCamelCase,  _UpperCamelCase  )
    elif hasattr(_UpperCamelCase,  '''Index'''  ) and isinstance(_UpperCamelCase,  ast.Index  ):
        return evaluate_ast(expression.value,  _UpperCamelCase,  _UpperCamelCase  )
    elif isinstance(_UpperCamelCase,  ast.JoinedStr  ):
        return "".join([str(evaluate_ast(_UpperCamelCase,  _UpperCamelCase,  _UpperCamelCase  )  ) for v in expression.values]  )
    elif isinstance(_UpperCamelCase,  ast.List  ):
        # List -> evaluate all elements
        return [evaluate_ast(_UpperCamelCase,  _UpperCamelCase,  _UpperCamelCase  ) for elt in expression.elts]
    elif isinstance(_UpperCamelCase,  ast.Name  ):
        # Name -> pick up the value in the state
        return evaluate_name(_UpperCamelCase,  _UpperCamelCase,  _UpperCamelCase  )
    elif isinstance(_UpperCamelCase,  ast.Subscript  ):
        # Subscript -> return the value of the indexing
        return evaluate_subscript(_UpperCamelCase,  _UpperCamelCase,  _UpperCamelCase  )
    else:
        # For now we refuse anything else. Let's add things as we need them.
        raise InterpretorError(F'''{expression.__class__.__name__} is not supported.'''  )





def 	_UpperCAmelCase       (				_UpperCamelCase       :      Dict,  _UpperCamelCase       :      int,  _UpperCamelCase       :      Dict  )      -> Union[str, Any]:
    A_      =  assign.targets
    A_      =  evaluate_ast(assign.value,  _UpperCamelCase,  _UpperCamelCase  )

    if len(_UpperCamelCase  ) == 1:
        A_      =  result
    else:
        if len(_UpperCamelCase  ) != len(_UpperCamelCase  ):
            raise InterpretorError(F'''Expected {len(_UpperCamelCase  )} values but got {len(_UpperCamelCase  )}.'''  )
        for var_name, r in zip(_UpperCamelCase,  _UpperCamelCase  ):
            A_      =  r
    return result





def 	_UpperCAmelCase       (				_UpperCamelCase       :      List[str],  _UpperCamelCase       :      List[str],  _UpperCamelCase       :      List[str]  )      -> Any:
    if not isinstance(call.func,  ast.Name  ):
        raise InterpretorError(
            F'''It is not permitted to evaluate other functions than the provided tools (tried to execute {call.func} of '''
            F'''type {type(call.func  )}.'''  )
    A_      =  call.func.id
    if func_name not in tools:
        raise InterpretorError(
            F'''It is not permitted to evaluate other functions than the provided tools (tried to execute {call.func.id}).'''  )

    A_      =  tools[func_name]
    # Todo deal with args
    A_      =  [evaluate_ast(_UpperCamelCase,  _UpperCamelCase,  _UpperCamelCase  ) for arg in call.args]
    A_      =  {keyword.arg: evaluate_ast(keyword.value,  _UpperCamelCase,  _UpperCamelCase  ) for keyword in call.keywords}
    return func(*_UpperCamelCase,  **_UpperCamelCase  )





def 	_UpperCAmelCase       (				_UpperCamelCase       :      Tuple,  _UpperCamelCase       :      Optional[Any],  _UpperCamelCase       :      Optional[int]  )      -> List[str]:
    A_      =  evaluate_ast(subscript.slice,  _UpperCamelCase,  _UpperCamelCase  )
    A_      =  evaluate_ast(subscript.value,  _UpperCamelCase,  _UpperCamelCase  )
    if isinstance(_UpperCamelCase,  (list, tuple)  ):
        return value[int(_UpperCamelCase  )]
    if index in value:
        return value[index]
    if isinstance(_UpperCamelCase,  _UpperCamelCase  ) and isinstance(_UpperCamelCase,  _UpperCamelCase  ):
        A_      =  difflib.get_close_matches(_UpperCamelCase,  list(value.keys()  )  )
        if len(_UpperCamelCase  ) > 0:
            return value[close_matches[0]]

    raise InterpretorError(F'''Could not index {value} with \'{index}\'.'''  )





def 	_UpperCAmelCase       (				_UpperCamelCase       :      str,  _UpperCamelCase       :      str,  _UpperCamelCase       :      int  )      -> Dict:
    if name.id in state:
        return state[name.id]
    A_      =  difflib.get_close_matches(name.id,  list(state.keys()  )  )
    if len(_UpperCamelCase  ) > 0:
        return state[close_matches[0]]
    raise InterpretorError(F'''The variable `{name.id}` is not defined.'''  )





def 	_UpperCAmelCase       (				_UpperCamelCase       :      Dict,  _UpperCamelCase       :      Optional[Any],  _UpperCamelCase       :      int  )      -> Any:
    if len(condition.ops  ) > 1:
        raise InterpretorError('''Cannot evaluate conditions with multiple operators'''  )

    A_      =  evaluate_ast(condition.left,  _UpperCamelCase,  _UpperCamelCase  )
    A_      =  condition.ops[0]
    A_      =  evaluate_ast(condition.comparators[0],  _UpperCamelCase,  _UpperCamelCase  )

    if isinstance(_UpperCamelCase,  ast.Eq  ):
        return left == right
    elif isinstance(_UpperCamelCase,  ast.NotEq  ):
        return left != right
    elif isinstance(_UpperCamelCase,  ast.Lt  ):
        return left < right
    elif isinstance(_UpperCamelCase,  ast.LtE  ):
        return left <= right
    elif isinstance(_UpperCamelCase,  ast.Gt  ):
        return left > right
    elif isinstance(_UpperCamelCase,  ast.GtE  ):
        return left >= right
    elif isinstance(_UpperCamelCase,  ast.Is  ):
        return left is right
    elif isinstance(_UpperCamelCase,  ast.IsNot  ):
        return left is not right
    elif isinstance(_UpperCamelCase,  ast.In  ):
        return left in right
    elif isinstance(_UpperCamelCase,  ast.NotIn  ):
        return left not in right
    else:
        raise InterpretorError(F'''Operator not supported: {comparator}'''  )





def 	_UpperCAmelCase       (				_UpperCamelCase       :      Optional[int],  _UpperCamelCase       :      List[Any],  _UpperCamelCase       :      List[str]  )      -> Optional[Any]:
    A_      =  None
    if evaluate_condition(if_statement.test,  _UpperCamelCase,  _UpperCamelCase  ):
        for line in if_statement.body:
            A_      =  evaluate_ast(_UpperCamelCase,  _UpperCamelCase,  _UpperCamelCase  )
            if line_result is not None:
                A_      =  line_result
    else:
        for line in if_statement.orelse:
            A_      =  evaluate_ast(_UpperCamelCase,  _UpperCamelCase,  _UpperCamelCase  )
            if line_result is not None:
                A_      =  line_result
    return result





def 	_UpperCAmelCase       (				_UpperCamelCase       :      List[Any],  _UpperCamelCase       :      Optional[Any],  _UpperCamelCase       :      str  )      -> Dict:
    A_      =  None
    A_      =  evaluate_ast(for_loop.iter,  _UpperCamelCase,  _UpperCamelCase  )
    for counter in iterator:
        A_      =  counter
        for expression in for_loop.body:
            A_      =  evaluate_ast(_UpperCamelCase,  _UpperCamelCase,  _UpperCamelCase  )
            if line_result is not None:
                A_      =  line_result
    return result







'''simple docstring'''




import json
import os
import tempfile

import datasets
from utils import generate_example_dataset, get_duration


__snake_case  : Dict		       =    500_000

__snake_case       ,       __snake_case  : List[str]		       =    os.path.split(__file__)
__snake_case  : Optional[Any]		       =    os.path.join(RESULTS_BASEPATH, 'results', RESULTS_FILENAME.replace('.py', '.json'))





@get_duration
def 	_UpperCAmelCase       (				_UpperCamelCase       :      datasets.Dataset  )      -> Union[str, Any]:
    A_      =  dataset.select(range(0,  len(_UpperCamelCase  ),  2  )  )





@get_duration
def 	_UpperCAmelCase       (				_UpperCamelCase       :      datasets.Dataset  )      -> str:
    A_      =  dataset.sort('''numbers'''  )





@get_duration
def 	_UpperCAmelCase       (				_UpperCamelCase       :      datasets.Dataset  )      -> str:
    A_      =  dataset.shuffle()





@get_duration
def 	_UpperCAmelCase       (				_UpperCamelCase       :      datasets.Dataset  )      -> Optional[Any]:
    A_      =  dataset.train_test_split(0.1  )





@get_duration
def 	_UpperCAmelCase       (				_UpperCamelCase       :      datasets.Dataset,  _UpperCamelCase       :      Optional[int]=10  )      -> Dict:
    for shard_id in range(_UpperCamelCase  ):
        A_      =  dataset.shard(_UpperCamelCase,  _UpperCamelCase  )





def 	_UpperCAmelCase       (				)      -> Any:
    A_      =  {'''num examples''': SPEED_TEST_N_EXAMPLES}
    A_      =  (select, sort, shuffle, train_test_split, shard)
    with tempfile.TemporaryDirectory() as tmp_dir:
        print('''generating dataset'''  )
        A_      =  datasets.Features({'''text''': datasets.Value('''string'''  ), '''numbers''': datasets.Value('''float32'''  )}  )
        A_      =  generate_example_dataset(
            os.path.join(_UpperCamelCase,  '''dataset.arrow'''  ),  _UpperCamelCase,  num_examples=_UpperCamelCase  )
        print('''Functions'''  )
        for func in functions:
            print(func.__name__  )
            A_      =  func(_UpperCamelCase  )

    with open(_UpperCamelCase,  '''wb'''  ) as f:
        f.write(json.dumps(_UpperCamelCase  ).encode('''utf-8'''  )  )


if __name__ == "__main__":  # useful to run the profiler
  benchmark_indices_mapping()






