



# 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 	snake_case_			(lowercase__	):




       """simple docstring"""






       pass







def  _snake_case			(							__snake_case       ,	__snake_case       ,	__snake_case=None       ,	__snake_case=False     )  ->  Optional[int]:






      '''simple docstring'''







      try:
            UpperCAmelCase_	:						Optional[int]					   =		ast.parse(__snake_case     )
      except SyntaxError as e:
            print("The code generated by the agent is not valid.\n"       ,	__snake_case     )
            return
      if state is None:
            UpperCAmelCase_	:						Dict					   =		{}
      UpperCAmelCase_	:						Union[str, Any]					   =		None
      for idx, node in enumerate(expression.body     ):
            try:
                  UpperCAmelCase_	:						Dict					   =		evaluate_ast(__snake_case       ,	__snake_case       ,	__snake_case     )
            except InterpretorError as e:
                  UpperCAmelCase_	:						Tuple					   =		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(__snake_case     )
                  break
            if line_result is not None:
                  UpperCAmelCase_	:						Tuple					   =		line_result

      return result







def  _snake_case			(							__snake_case       ,	__snake_case       ,	__snake_case     )  ->  Optional[int]:






      '''simple docstring'''







      if isinstance(__snake_case       ,	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(__snake_case       ,	__snake_case       ,	__snake_case     )
      elif isinstance(__snake_case       ,	ast.Call     ):
            # Function call -> we return the value of the function call
            return evaluate_call(__snake_case       ,	__snake_case       ,	__snake_case     )
      elif isinstance(__snake_case       ,	ast.Constant     ):
            # Constant -> just return the value
            return expression.value
      elif isinstance(__snake_case       ,	ast.Dict     ):
            # Dict -> evaluate all keys and values
            UpperCAmelCase_	:						int					   =		[evaluate_ast(__snake_case       ,	__snake_case       ,	__snake_case     ) for k in expression.keys]
            UpperCAmelCase_	:						Union[str, Any]					   =		[evaluate_ast(__snake_case       ,	__snake_case       ,	__snake_case     ) for v in expression.values]
            return dict(zip(__snake_case       ,	__snake_case     )     )
      elif isinstance(__snake_case       ,	ast.Expr     ):
            # Expression -> evaluate the content
            return evaluate_ast(expression.value       ,	__snake_case       ,	__snake_case     )
      elif isinstance(__snake_case       ,	ast.For     ):
            # For loop -> execute the loop
            return evaluate_for(__snake_case       ,	__snake_case       ,	__snake_case     )
      elif isinstance(__snake_case       ,	ast.FormattedValue     ):
            # Formatted value (part of f-string) -> evaluate the content and return
            return evaluate_ast(expression.value       ,	__snake_case       ,	__snake_case     )
      elif isinstance(__snake_case       ,	ast.If     ):
            # If -> execute the right branch
            return evaluate_if(__snake_case       ,	__snake_case       ,	__snake_case     )
      elif hasattr(__snake_case       ,	"Index"     ) and isinstance(__snake_case       ,	ast.Index     ):
            return evaluate_ast(expression.value       ,	__snake_case       ,	__snake_case     )
      elif isinstance(__snake_case       ,	ast.JoinedStr     ):
            return "".join([str(evaluate_ast(__snake_case       ,	__snake_case       ,	__snake_case     )     ) for v in expression.values]     )
      elif isinstance(__snake_case       ,	ast.List     ):
            # List -> evaluate all elements
            return [evaluate_ast(__snake_case       ,	__snake_case       ,	__snake_case     ) for elt in expression.elts]
      elif isinstance(__snake_case       ,	ast.Name     ):
            # Name -> pick up the value in the state
            return evaluate_name(__snake_case       ,	__snake_case       ,	__snake_case     )
      elif isinstance(__snake_case       ,	ast.Subscript     ):
            # Subscript -> return the value of the indexing
            return evaluate_subscript(__snake_case       ,	__snake_case       ,	__snake_case     )
      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  _snake_case			(							__snake_case       ,	__snake_case       ,	__snake_case     )  ->  Tuple:






      '''simple docstring'''







      UpperCAmelCase_	:						Union[str, Any]					   =		assign.targets
      UpperCAmelCase_	:						Optional[int]					   =		evaluate_ast(assign.value       ,	__snake_case       ,	__snake_case     )

      if len(__snake_case     ) == 1:
            UpperCAmelCase_	:						Tuple					   =		result
      else:
            if len(__snake_case     ) != len(__snake_case     ):
                  raise InterpretorError(F"""Expected {len(__snake_case     )} values but got {len(__snake_case     )}."""     )
            for var_name, r in zip(__snake_case       ,	__snake_case     ):
                  UpperCAmelCase_	:						str					   =		r
      return result







def  _snake_case			(							__snake_case       ,	__snake_case       ,	__snake_case     )  ->  List[Any]:






      '''simple docstring'''







      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     )}."""     )
      UpperCAmelCase_	:						str					   =		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})."""     )

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







def  _snake_case			(							__snake_case       ,	__snake_case       ,	__snake_case     )  ->  Optional[int]:






      '''simple docstring'''







      UpperCAmelCase_	:						Optional[Any]					   =		evaluate_ast(subscript.slice       ,	__snake_case       ,	__snake_case     )
      UpperCAmelCase_	:						Tuple					   =		evaluate_ast(subscript.value       ,	__snake_case       ,	__snake_case     )
      if isinstance(__snake_case       ,	(list, tuple)     ):
            return value[int(__snake_case     )]
      if index in value:
            return value[index]
      if isinstance(__snake_case       ,	__snake_case     ) and isinstance(__snake_case       ,	__snake_case     ):
            UpperCAmelCase_	:						Optional[int]					   =		difflib.get_close_matches(__snake_case       ,	list(value.keys()     )     )
            if len(__snake_case     ) > 0:
                  return value[close_matches[0]]

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







def  _snake_case			(							__snake_case       ,	__snake_case       ,	__snake_case     )  ->  List[Any]:






      '''simple docstring'''







      if name.id in state:
            return state[name.id]
      UpperCAmelCase_	:						Any					   =		difflib.get_close_matches(name.id       ,	list(state.keys()     )     )
      if len(__snake_case     ) > 0:
            return state[close_matches[0]]
      raise InterpretorError(F"""The variable `{name.id}` is not defined."""     )







def  _snake_case			(							__snake_case       ,	__snake_case       ,	__snake_case     )  ->  int:






      '''simple docstring'''







      if len(condition.ops     ) > 1:
            raise InterpretorError("Cannot evaluate conditions with multiple operators"     )

      UpperCAmelCase_	:						Optional[Any]					   =		evaluate_ast(condition.left       ,	__snake_case       ,	__snake_case     )
      UpperCAmelCase_	:						Any					   =		condition.ops[0]
      UpperCAmelCase_	:						Optional[int]					   =		evaluate_ast(condition.comparators[0]       ,	__snake_case       ,	__snake_case     )

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







def  _snake_case			(							__snake_case       ,	__snake_case       ,	__snake_case     )  ->  Dict:






      '''simple docstring'''







      UpperCAmelCase_	:						int					   =		None
      if evaluate_condition(if_statement.test       ,	__snake_case       ,	__snake_case     ):
            for line in if_statement.body:
                  UpperCAmelCase_	:						List[Any]					   =		evaluate_ast(__snake_case       ,	__snake_case       ,	__snake_case     )
                  if line_result is not None:
                        UpperCAmelCase_	:						Union[str, Any]					   =		line_result
      else:
            for line in if_statement.orelse:
                  UpperCAmelCase_	:						Any					   =		evaluate_ast(__snake_case       ,	__snake_case       ,	__snake_case     )
                  if line_result is not None:
                        UpperCAmelCase_	:						Dict					   =		line_result
      return result







def  _snake_case			(							__snake_case       ,	__snake_case       ,	__snake_case     )  ->  Optional[Any]:






      '''simple docstring'''







      UpperCAmelCase_	:						int					   =		None
      UpperCAmelCase_	:						Any					   =		evaluate_ast(for_loop.iter       ,	__snake_case       ,	__snake_case     )
      for counter in iterator:
            UpperCAmelCase_	:						str					   =		counter
            for expression in for_loop.body:
                  UpperCAmelCase_	:						List[str]					   =		evaluate_ast(__snake_case       ,	__snake_case       ,	__snake_case     )
                  if line_result is not None:
                        UpperCAmelCase_	:						Tuple					   =		line_result
      return result





import json
import os
from collections.abc import Mapping
from typing import Dict, List, Optional, Union

import numpy as np
import torch

from .imports import is_safetensors_available







def  _snake_case			(							__snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case=None     )  ->  Any:






      '''simple docstring'''







      UpperCAmelCase_	:						Dict					   =		None
      # Check the string instead of the dtype to be compatible with versions of PyTorch that don't have bfloat16.
      if str(weight.dtype     ) == "torch.bfloat16":
            # Need to reinterpret the underlined data as int16 since NumPy does not handle bfloat16s.
            UpperCAmelCase_	:						Optional[Any]					   =		weight.view(torch.intaa     )
            UpperCAmelCase_	:						Dict					   =		"bfloat16"
      UpperCAmelCase_	:						Optional[int]					   =		weight.cpu().numpy()
      UpperCAmelCase_	:						Tuple					   =		os.path.join(__snake_case       ,	F"""{weight_name}.dat"""     )
      if index is not None:
            if dtype is None:
                  UpperCAmelCase_	:						int					   =		str(array.dtype     )
            UpperCAmelCase_	:						Any					   =		{"dtype": dtype, "shape": list(array.shape     )}
      if array.ndim == 0:
            UpperCAmelCase_	:						List[str]					   =		array[None]
      UpperCAmelCase_	:						List[Any]					   =		np.memmap(__snake_case       ,	dtype=array.dtype       ,	mode="w+"       ,	shape=array.shape     )
      UpperCAmelCase_	:						List[str]					   =		array[:]
      file_array.flush()
      return index







def  _snake_case			(							__snake_case       ,	__snake_case     )  ->  Any:






      '''simple docstring'''







      UpperCAmelCase_	:						Union[str, Any]					   =		tuple(weight_info["shape"]     )
      if shape == ():
            # NumPy memory-mapped arrays can't have 0 dims so it was saved as 1d tensor
            UpperCAmelCase_	:						Dict					   =		(1,)

      UpperCAmelCase_	:						Tuple					   =		weight_info["dtype"]
      if dtype == "bfloat16":
            # NumPy does not support bfloat16 so this was saved as a int16
            UpperCAmelCase_	:						Tuple					   =		"int16"

      UpperCAmelCase_	:						Union[str, Any]					   =		np.memmap(__snake_case       ,	dtype=__snake_case       ,	shape=__snake_case       ,	mode="r"     )

      if len(weight_info["shape"]     ) == 0:
            UpperCAmelCase_	:						Optional[Any]					   =		weight[0]
      UpperCAmelCase_	:						List[str]					   =		torch.tensor(__snake_case     )
      if weight_info["dtype"] == "bfloat16":
            UpperCAmelCase_	:						Union[str, Any]					   =		weight.view(torch.bfloataa     )

      return weight







def  _snake_case			(							__snake_case       ,	__snake_case     )  ->  List[Any]:






      '''simple docstring'''







      if index is None or len(__snake_case     ) == 0:
            # Nothing to save
            return

      UpperCAmelCase_	:						str					   =		os.path.join(__snake_case       ,	"index.json"     )
      if os.path.isfile(__snake_case     ):
            with open(__snake_case       ,	"r"       ,	encoding="utf-8"     ) as f:
                  UpperCAmelCase_	:						str					   =		json.load(__snake_case     )
      else:
            UpperCAmelCase_	:						Dict					   =		{}
      current_index.update(__snake_case     )

      with open(__snake_case       ,	"w"       ,	encoding="utf-8"     ) as f:
            json.dump(__snake_case       ,	__snake_case       ,	indent=2     )







def  _snake_case			(							__snake_case       ,	__snake_case     )  ->  Union[str, Any]:






      '''simple docstring'''







      os.makedirs(__snake_case       ,	exist_ok=__snake_case     )
      UpperCAmelCase_	:						Any					   =		{}
      for name, parameter in state_dict.items():
            UpperCAmelCase_	:						Tuple					   =		offload_weight(__snake_case       ,	__snake_case       ,	__snake_case       ,	index=__snake_case     )

      # Update index
      save_offload_index(__snake_case       ,	__snake_case     )


class 	snake_case_			(lowercase__	):




       """simple docstring"""






       def __init__(  self  ,lowercase  ,lowercase):

             """simple docstring"""





             UpperCAmelCase_	:						Optional[Any]					   =		dataset
             UpperCAmelCase_	:						Any					   =		prefix



       def __getitem__(  self  ,lowercase):

             """simple docstring"""





             return self.dataset[F"""{self.prefix}{key}"""]



       def __iter__(  self):

             """simple docstring"""





             return iter([key for key in self.dataset if key.startswith(self.prefix)])



       def __len__(  self):

             """simple docstring"""





             return len(self.dataset)


class 	snake_case_			(lowercase__	):




       """simple docstring"""






       def __init__(  self  ,lowercase = None  ,lowercase = None  ,lowercase = None  ,lowercase=None  ,):

             """simple docstring"""





             if state_dict is None and save_folder is None:
                   raise ValueError("Need either a `state_dict` or a `save_folder` containing offloaded weights.")

             UpperCAmelCase_	:						Tuple					   =		{} if state_dict is None else state_dict
             UpperCAmelCase_	:						List[Any]					   =		save_folder
             if index is None and save_folder is not None:
                   with open(os.path.join(lowercase  ,"index.json")) as f:
                         UpperCAmelCase_	:						Dict					   =		json.load(lowercase)
             UpperCAmelCase_	:						List[str]					   =		{} if index is None else index
             UpperCAmelCase_	:						int					   =		list(self.state_dict.keys())
             self.all_keys.extend([key for key in self.index if key not in self.all_keys])
             UpperCAmelCase_	:						Optional[Any]					   =		device



       def __getitem__(  self  ,lowercase):

             """simple docstring"""





             if key in self.state_dict:
                   return self.state_dict[key]
             UpperCAmelCase_	:						str					   =		self.index[key]
             if weight_info.get("safetensors_file") is not None:
                   if not is_safetensors_available():
                         raise ImportError("These offloaded weights require the use of safetensors: `pip install safetensors`.")

                   from safetensors import safe_open

                   UpperCAmelCase_	:						Any					   =		"cpu" if self.device is None else self.device
                   with safe_open(weight_info["safetensors_file"]  ,framework="pt"  ,device=lowercase) as f:
                         UpperCAmelCase_	:						int					   =		f.get_tensor(weight_info.get("weight_name"  ,lowercase))

                   if "dtype" in weight_info:
                         return tensor.to(getattr(lowercase  ,weight_info["dtype"]))
                   else:
                         return tensor

             UpperCAmelCase_	:						Optional[int]					   =		os.path.join(self.save_folder  ,F"""{key}.dat""")
             return load_offloaded_weight(lowercase  ,lowercase)



       def __iter__(  self):

             """simple docstring"""





             return iter(self.all_keys)



       def __len__(  self):

             """simple docstring"""





             return len(self.all_keys)







def  _snake_case			(							__snake_case       ,	__snake_case     )  ->  Optional[Any]:






      '''simple docstring'''







      UpperCAmelCase_	:						str					   =		{}
      for module_name in submodule_names:
            # We want to catch module_name parameter (module_name.xxx) or potentially module_name, but not any of the
            # submodules that could being like module_name (transformers.h.1 and transformers.h.10 for instance)
            result.update(
                {
                    key: param
                    for key, param in state_dict.items()
                    if key == module_name or key.startswith(module_name + "."     )
                }     )
      return result
