




import re

import jax.numpy as jnp
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.random import PRNGKey

from ..utils import logging


lowerCamelCase	:	int				=   logging.get_logger(__name__)




def 						SCREAMING_SNAKE_CASE__ (   lowercase					)  ->       Optional[int]:
   snake_case				:    Dict           =	r"\w+[.]\d+"
   snake_case				:    List[Any]           =	re.findall(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE					)
   for pat in pats:
      snake_case				:    List[Any]           =	key.replace(_SCREAMING_SNAKE_CASE						,"""_""".join(pat.split("""."""					)					)					)
   return key




def 						SCREAMING_SNAKE_CASE__ (   lowercase						,lowercase						,lowercase					)  ->       Optional[Any]:
   snake_case				:    List[Any]           =	pt_tuple_key[:-1] + ("scale",)
   if (
       any("""norm""" in str_ for str_ in pt_tuple_key					)
       and (pt_tuple_key[-1] == "bias")
       and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)
       and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
   ):
      snake_case				:    Optional[int]           =	pt_tuple_key[:-1] + ("scale",)
      return renamed_pt_tuple_key, pt_tensor
   elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
      snake_case				:    Optional[int]           =	pt_tuple_key[:-1] + ("scale",)
      return renamed_pt_tuple_key, pt_tensor

   # embedding
   if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
      snake_case				:    List[Any]           =	pt_tuple_key[:-1] + ("embedding",)
      return renamed_pt_tuple_key, pt_tensor

   # conv layer
   snake_case				:    List[str]           =	pt_tuple_key[:-1] + ("kernel",)
   if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:
      snake_case				:    List[str]           =	pt_tensor.transpose(2						,3						,1						,0					)
      return renamed_pt_tuple_key, pt_tensor

   # linear layer
   snake_case				:    Tuple           =	pt_tuple_key[:-1] + ("kernel",)
   if pt_tuple_key[-1] == "weight":
      snake_case				:    str           =	pt_tensor.T
      return renamed_pt_tuple_key, pt_tensor

   # old PyTorch layer norm weight
   snake_case				:    int           =	pt_tuple_key[:-1] + ("weight",)
   if pt_tuple_key[-1] == "gamma":
      return renamed_pt_tuple_key, pt_tensor

   # old PyTorch layer norm bias
   snake_case				:    Union[str, Any]           =	pt_tuple_key[:-1] + ("bias",)
   if pt_tuple_key[-1] == "beta":
      return renamed_pt_tuple_key, pt_tensor

   return pt_tuple_key, pt_tensor




def 						SCREAMING_SNAKE_CASE__ (   lowercase						,lowercase						,lowercase=42					)  ->       Dict:
   snake_case				:    Union[str, Any]           =	{k: v.numpy() for k, v in pt_state_dict.items()}

   # Step 2: Since the model is stateless, get random Flax params
   snake_case				:    Tuple           =	flax_model.init_weights(PRNGKey(_SCREAMING_SNAKE_CASE					)					)

   snake_case				:    Optional[int]           =	flatten_dict(_SCREAMING_SNAKE_CASE					)
   snake_case				:    Optional[Any]           =	{}

   # Need to change some parameters name to match Flax names
   for pt_key, pt_tensor in pt_state_dict.items():
      snake_case				:    str           =	rename_key(_SCREAMING_SNAKE_CASE					)
      snake_case				:    int           =	tuple(renamed_pt_key.split("""."""					)					)

      # Correctly rename weight parameters
      snake_case				:    int           =	rename_key_and_reshape_tensor(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE					)

      if flax_key in random_flax_state_dict:
         if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
            raise ValueError(
                f"""PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape """
                f"""{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}."""					)

        # also add unexpected weight so that warning is thrown
      snake_case				:    Optional[int]           =	jnp.asarray(_SCREAMING_SNAKE_CASE					)

   return unflatten_dict(_SCREAMING_SNAKE_CASE					)











import pytest
import requests

from datasets.utils.file_utils import http_head

from .utils import OfflineSimulationMode, RequestWouldHangIndefinitelyError, offline




@pytest.mark.integration
def 						SCREAMING_SNAKE_CASE__ (   )  ->       List[Any]:
   with offline(OfflineSimulationMode.CONNECTION_TIMES_OUT					):
      with pytest.raises(lowercase					):
         requests.request("""GET"""						,"""https://huggingface.co"""					)
      with pytest.raises(requests.exceptions.ConnectTimeout					):
         requests.request("""GET"""						,"""https://huggingface.co"""						,timeout=1.0					)




@pytest.mark.integration
def 						SCREAMING_SNAKE_CASE__ (   )  ->       Any:
   with offline(OfflineSimulationMode.CONNECTION_FAILS					):
      with pytest.raises(requests.exceptions.ConnectionError					):
         requests.request("""GET"""						,"""https://huggingface.co"""					)




def 						SCREAMING_SNAKE_CASE__ (   )  ->       int:
   with offline(OfflineSimulationMode.HF_DATASETS_OFFLINE_SET_TO_1					):
      with pytest.raises(lowercase					):
         http_head("""https://huggingface.co"""					)





