





import json
import os
import pickle
import subprocess
from hashlib import mda
from pathlib import Path
from tempfile import gettempdir
from textwrap import dedent
from types import CodeType, FunctionType
from unittest import TestCase
from unittest.mock import patch

import pytest
from multiprocess import Pool
from packaging import version

import datasets
from datasets.fingerprint import Hasher, fingerprint_transform
from datasets.table import InMemoryTable

from .utils import (
    require_regex,
    require_spacy,
    require_spacy_model,
    require_tiktoken,
    require_torch,
    require_transformers,
)







class        __snake_case			:





       '''simple docstring'''







       def __init__(    self					:       Optional[Any]				,							A					:       Tuple						):
         __snake_case:		Optional[Any]					    =		foo






       def __call__(    self					:       Optional[Any]						):
         return self.foo







class        __snake_case			(				datasets.Dataset     ):





       '''simple docstring'''







       @fingerprint_transform(inplace=_lowerCAmelCase						)
       def 				UpperCAmelCase__					(    self					:       List[str]				,							A					:       List[str]				,							*A					:       List[Any]				,							**A					:       Optional[Any]						):
         return DatasetChild(self.data				,							fingerprint=_lowerCAmelCase						)






       @fingerprint_transform(inplace=_lowerCAmelCase						)
       def 				UpperCAmelCase__					(    self					:       Any				,							A					:       List[Any]				,							*A					:       List[str]				,							**A					:       Optional[Any]						):
         return DatasetChild(self.data				,							fingerprint=_lowerCAmelCase						)







class        __snake_case			:





       '''simple docstring'''







       def __init__(    self					:       int				,							A					:       int						):
         __snake_case:		List[Any]					    =		callable




       def __call__(    self					:       Optional[Any]				,							*A					:       Dict				,							**A					:       Optional[Any]						):
         if self.callable is not None:
           return self.callable(*_lowerCAmelCase				,							**_lowerCAmelCase						)






       def __getstate__(    self					:       int						):
         raise pickle.PicklingError()







class        __snake_case			(				__lowerCamelCase     ):





       '''simple docstring'''







       @require_transformers
       @pytest.mark.integration
       def 				UpperCAmelCase__					(    self					:       int						):
         from transformers import AutoTokenizer

         def encode(A					:       Tuple						):
           return tokenizer(_lowerCAmelCase						)

         # TODO: add hash consistency tests across sessions
         __snake_case:		Optional[Any]					    =		AutoTokenizer.from_pretrained("""bert-base-uncased"""						)
         __snake_case:		Tuple					    =		mda(datasets.utils.py_utils.dumps(_lowerCAmelCase						)						).hexdigest()
         __snake_case:		List[Any]					    =		mda(datasets.utils.py_utils.dumps(lambda A						: tokenizer(_lowerCAmelCase						)						)						).hexdigest()
         __snake_case:		List[str]					    =		mda(datasets.utils.py_utils.dumps(_lowerCAmelCase						)						).hexdigest()
         __snake_case:		List[str]					    =		AutoTokenizer.from_pretrained("""bert-base-cased"""						)
         __snake_case:		int					    =		mda(datasets.utils.py_utils.dumps(_lowerCAmelCase						)						).hexdigest()
         __snake_case:		List[str]					    =		mda(datasets.utils.py_utils.dumps(lambda A						: tokenizer(_lowerCAmelCase						)						)						).hexdigest()
         __snake_case:		int					    =		mda(datasets.utils.py_utils.dumps(_lowerCAmelCase						)						).hexdigest()
         __snake_case:		str					    =		AutoTokenizer.from_pretrained("""bert-base-uncased"""						)
         __snake_case:		Dict					    =		mda(datasets.utils.py_utils.dumps(_lowerCAmelCase						)						).hexdigest()
         __snake_case:		Optional[int]					    =		mda(datasets.utils.py_utils.dumps(lambda A						: tokenizer(_lowerCAmelCase						)						)						).hexdigest()
         __snake_case:		Dict					    =		mda(datasets.utils.py_utils.dumps(_lowerCAmelCase						)						).hexdigest()
         self.assertEqual(_lowerCAmelCase				,							_lowerCAmelCase						)
         self.assertNotEqual(_lowerCAmelCase				,							_lowerCAmelCase						)
         self.assertEqual(_lowerCAmelCase				,							_lowerCAmelCase						)
         self.assertNotEqual(_lowerCAmelCase				,							_lowerCAmelCase						)
         self.assertEqual(_lowerCAmelCase				,							_lowerCAmelCase						)
         self.assertNotEqual(_lowerCAmelCase				,							_lowerCAmelCase						)




       @require_transformers
       @pytest.mark.integration
       def 				UpperCAmelCase__					(    self					:       int						):
         from transformers import AutoTokenizer

         __snake_case:		Dict					    =		AutoTokenizer.from_pretrained("""gpt2"""						)
         __snake_case:		Optional[Any]					    =		mda(datasets.utils.py_utils.dumps(_lowerCAmelCase						)						).hexdigest()
         tokenizer("""Hello world !"""						)  # call once to change the tokenizer's cache
         __snake_case:		Optional[int]					    =		mda(datasets.utils.py_utils.dumps(_lowerCAmelCase						)						).hexdigest()
         self.assertEqual(_lowerCAmelCase				,							_lowerCAmelCase						)






       @require_regex
       def 				UpperCAmelCase__					(    self					:       int						):
         import regex

         __snake_case:		Optional[int]					    =		regex.Regex("""foo"""						)
         __snake_case:		Union[str, Any]					    =		mda(datasets.utils.py_utils.dumps(_lowerCAmelCase						)						).hexdigest()
         __snake_case:		List[Any]					    =		regex.Regex("""bar"""						)
         __snake_case:		List[Any]					    =		mda(datasets.utils.py_utils.dumps(_lowerCAmelCase						)						).hexdigest()
         __snake_case:		Tuple					    =		regex.Regex("""foo"""						)
         __snake_case:		Tuple					    =		mda(datasets.utils.py_utils.dumps(_lowerCAmelCase						)						).hexdigest()
         self.assertEqual(_lowerCAmelCase				,							_lowerCAmelCase						)
         self.assertNotEqual(_lowerCAmelCase				,							_lowerCAmelCase						)







class        __snake_case			(				__lowerCamelCase     ):





       '''simple docstring'''







       def 				UpperCAmelCase__					(    self					:       str						):
         def func():
           return foo

         __snake_case:		Optional[Any]					    =		[0]
         __snake_case:		Union[str, Any]					    =		mda(datasets.utils.py_utils.dumps(_lowerCAmelCase						)						).hexdigest()
         __snake_case:		int					    =		[1]
         __snake_case:		Dict					    =		mda(datasets.utils.py_utils.dumps(_lowerCAmelCase						)						).hexdigest()
         __snake_case:		Optional[Any]					    =		[0]
         __snake_case:		Tuple					    =		mda(datasets.utils.py_utils.dumps(_lowerCAmelCase						)						).hexdigest()
         self.assertEqual(_lowerCAmelCase				,							_lowerCAmelCase						)
         self.assertNotEqual(_lowerCAmelCase				,							_lowerCAmelCase						)




       def 				UpperCAmelCase__					(    self					:       Union[str, Any]						):
         def func():
           pass

         __snake_case:		int					    =		mda(datasets.utils.py_utils.dumps(_lowerCAmelCase						)						).hexdigest()

         def func():
           pass

         __snake_case:		str					    =		mda(datasets.utils.py_utils.dumps(_lowerCAmelCase						)						).hexdigest()
         self.assertEqual(_lowerCAmelCase				,							_lowerCAmelCase						)




       def 				UpperCAmelCase__					(    self					:       Dict						):
         __snake_case:		Union[str, Any]					    =		mda(datasets.utils.py_utils.dumps(Foo([0]						)						)						).hexdigest()
         __snake_case:		Optional[int]					    =		mda(datasets.utils.py_utils.dumps(Foo([1]						)						)						).hexdigest()
         __snake_case:		Any					    =		mda(datasets.utils.py_utils.dumps(Foo([0]						)						)						).hexdigest()
         self.assertEqual(_lowerCAmelCase				,							_lowerCAmelCase						)
         self.assertNotEqual(_lowerCAmelCase				,							_lowerCAmelCase						)




       def 				UpperCAmelCase__					(    self					:       str						):
         __snake_case:		int					    =		mda(datasets.utils.py_utils.dumps(Foo([0]						).__call__						)						).hexdigest()
         __snake_case:		List[str]					    =		mda(datasets.utils.py_utils.dumps(Foo([1]						).__call__						)						).hexdigest()
         __snake_case:		Union[str, Any]					    =		mda(datasets.utils.py_utils.dumps(Foo([0]						).__call__						)						).hexdigest()
         self.assertEqual(_lowerCAmelCase				,							_lowerCAmelCase						)
         self.assertNotEqual(_lowerCAmelCase				,							_lowerCAmelCase						)




       def 				UpperCAmelCase__					(    self					:       Union[str, Any]						):
         __snake_case:		Optional[int]					    =		(
             "co_argcount",
             "co_kwonlyargcount",
             "co_nlocals",
             "co_stacksize",
             "co_flags",
             "co_code",
             "co_consts",
             "co_names",
             "co_varnames",
             "co_filename",
             "co_name",
             "co_firstlineno",
             "co_lnotab",
             "co_freevars",
             "co_cellvars",
         )

         def _create_code_pyaa(*A					:       Optional[Any]						):
           if hasattr(args[-3]				,							"""encode"""						):
             __snake_case:		Any					    =		list(_lowerCAmelCase						)
             __snake_case:		str					    =		args[-3].encode()  # co_lnotab
             __snake_case:		List[Any]					    =		args[-10].encode()  # co_code
           if hasattr(_lowerCAmelCase				,							"""co_posonlyargcount"""						):
             if len(_lowerCAmelCase						) == 16:
               return CodeType(*_lowerCAmelCase						)
             elif len(_lowerCAmelCase						) == 15:
               return CodeType(args[0]				,							0				,							*args[1:]						)
             return CodeType(args[0]				,							0				,							0				,							*args[1:]						)
           elif hasattr(_lowerCAmelCase				,							"""co_kwonlyargcount"""						):
             if len(_lowerCAmelCase						) == 16:
               return CodeType(args[0]				,							*args[2:]						)
             elif len(_lowerCAmelCase						) == 15:
               return CodeType(*_lowerCAmelCase						)
             return CodeType(args[0]				,							0				,							*args[1:]						)
           if len(_lowerCAmelCase						) == 16:
             return CodeType(args[0]				,							*args[3:]						)
           elif len(_lowerCAmelCase						) == 15:
             return CodeType(args[0]				,							*args[2:]						)
           return CodeType(*_lowerCAmelCase						)

         def create_ipython_func(A					:       Dict				,							A					:       Dict						):
           def func():
             return returned_obj

           __snake_case:		Any					    =		func.__code__
           # Use _create_code from dill in order to make it work for different python versions
           if datasets.config.PY_VERSION < version.parse("""3.8"""						):
             __snake_case:		Union[str, Any]					    =		_create_code_pyaa(
                 *[getattr(_lowerCAmelCase				,							_lowerCAmelCase						) if k != """co_filename""" else co_filename for k in code_args_pyaa]						)
           else:
             __snake_case:		str					    =		code.replace(co_filename=_lowerCAmelCase						)
           return FunctionType(_lowerCAmelCase				,							func.__globals__				,							func.__name__				,							func.__defaults__				,							func.__closure__						)

         __snake_case:		Union[str, Any]					    =		"<ipython-input-2-e0383a102aae>", [0]
         __snake_case:		str					    =		mda(datasets.utils.py_utils.dumps(create_ipython_func(_lowerCAmelCase				,							_lowerCAmelCase						)						)						).hexdigest()
         __snake_case:		Any					    =		"<ipython-input-2-e0383a102aae>", [1]
         __snake_case:		Dict					    =		mda(datasets.utils.py_utils.dumps(create_ipython_func(_lowerCAmelCase				,							_lowerCAmelCase						)						)						).hexdigest()
         __snake_case:		List[Any]					    =		"<ipython-input-5-713f6613acf3>", [0]
         __snake_case:		str					    =		mda(datasets.utils.py_utils.dumps(create_ipython_func(_lowerCAmelCase				,							_lowerCAmelCase						)						)						).hexdigest()
         self.assertEqual(_lowerCAmelCase				,							_lowerCAmelCase						)
         self.assertNotEqual(_lowerCAmelCase				,							_lowerCAmelCase						)

         __snake_case:		Tuple					    =		os.path.join(gettempdir()				,							"""ipykernel_12345"""				,							"""321456789.py"""						), [0]
         __snake_case:		List[str]					    =		mda(datasets.utils.py_utils.dumps(create_ipython_func(_lowerCAmelCase				,							_lowerCAmelCase						)						)						).hexdigest()
         __snake_case:		Tuple					    =		os.path.join(gettempdir()				,							"""ipykernel_12345"""				,							"""321456789.py"""						), [1]
         __snake_case:		Dict					    =		mda(datasets.utils.py_utils.dumps(create_ipython_func(_lowerCAmelCase				,							_lowerCAmelCase						)						)						).hexdigest()
         __snake_case:		Union[str, Any]					    =		os.path.join(gettempdir()				,							"""ipykernel_12345"""				,							"""654123987.py"""						), [0]
         __snake_case:		Tuple					    =		mda(datasets.utils.py_utils.dumps(create_ipython_func(_lowerCAmelCase				,							_lowerCAmelCase						)						)						).hexdigest()
         self.assertEqual(_lowerCAmelCase				,							_lowerCAmelCase						)
         self.assertNotEqual(_lowerCAmelCase				,							_lowerCAmelCase						)






       def 				UpperCAmelCase__					(    self					:       Optional[Any]						):
         __snake_case:		List[Any]					    =		[0], [1]

         def func():
           return foo, bar

         __snake_case:		Tuple					    =		"__main__"

         def globalvars_mocka_side_effect(A					:       List[Any]				,							*A					:       Dict				,							**A					:       int						):
           return {"foo": foo, "bar": bar}

         def globalvars_mocka_side_effect(A					:       Dict				,							*A					:       List[Any]				,							**A					:       Any						):
           return {"bar": bar, "foo": foo}

         with patch("""dill.detect.globalvars"""				,							side_effect=_lowerCAmelCase						) as globalvars_mocka:
           __snake_case:		Union[str, Any]					    =		mda(datasets.utils.py_utils.dumps(_lowerCAmelCase						)						).hexdigest()
           self.assertGreater(globalvars_mocka.call_count				,							0						)
         with patch("""dill.detect.globalvars"""				,							side_effect=_lowerCAmelCase						) as globalvars_mocka:
           __snake_case:		int					    =		mda(datasets.utils.py_utils.dumps(_lowerCAmelCase						)						).hexdigest()
           self.assertGreater(globalvars_mocka.call_count				,							0						)
         self.assertEqual(_lowerCAmelCase				,							_lowerCAmelCase						)







class        __snake_case			(				__lowerCamelCase     ):





       '''simple docstring'''







       def 				UpperCAmelCase__					(    self					:       Any						):
         __snake_case:		Optional[int]					    =		Hasher.hash("""hello"""						)
         __snake_case:		str					    =		Hasher.hash("""hello"""						)
         __snake_case:		Tuple					    =		Hasher.hash("""there"""						)
         self.assertEqual(_lowerCAmelCase				,							_lowerCAmelCase						)
         self.assertNotEqual(_lowerCAmelCase				,							_lowerCAmelCase						)




       def 				UpperCAmelCase__					(    self					:       Optional[Any]						):
         __snake_case:		int					    =		Hasher.hash(Foo("""hello"""						)						)
         __snake_case:		List[str]					    =		Hasher.hash(Foo("""hello"""						)						)
         __snake_case:		int					    =		Hasher.hash(Foo("""there"""						)						)
         self.assertEqual(_lowerCAmelCase				,							_lowerCAmelCase						)
         self.assertNotEqual(_lowerCAmelCase				,							_lowerCAmelCase						)




       def 				UpperCAmelCase__					(    self					:       List[str]						):
         __snake_case:		Dict					    =		Hasher()
         for x in ["hello", Foo("""hello"""						)]:
           hasher.update(_lowerCAmelCase						)
         __snake_case:		List[Any]					    =		hasher.hexdigest()
         __snake_case:		Optional[int]					    =		Hasher()
         for x in ["hello", Foo("""hello"""						)]:
           hasher.update(_lowerCAmelCase						)
         __snake_case:		str					    =		hasher.hexdigest()
         __snake_case:		Optional[int]					    =		Hasher()
         for x in ["there", Foo("""there"""						)]:
           hasher.update(_lowerCAmelCase						)
         __snake_case:		Dict					    =		hasher.hexdigest()
         self.assertEqual(_lowerCAmelCase				,							_lowerCAmelCase						)
         self.assertNotEqual(_lowerCAmelCase				,							_lowerCAmelCase						)




       def 				UpperCAmelCase__					(    self					:       int						):
         with self.assertRaises(pickle.PicklingError						):
           Hasher.hash(UnpicklableCallable(Foo("""hello"""						)						)						)




       def 				UpperCAmelCase__					(    self					:       Tuple						):
         __snake_case:		str					    =		"abc"
         __snake_case:		Optional[int]					    =		[string, string]  # two strings have the same ids
         __snake_case:		Any					    =		[string, string]
         __snake_case:		str					    =		json.loads(f'''[\"{string}\", \"{string}\"]'''						)  # two strings have different ids
         self.assertIs(obja[0]				,							_lowerCAmelCase						)
         self.assertIs(obja[0]				,							obja[1]						)
         self.assertIs(obja[0]				,							_lowerCAmelCase						)
         self.assertIs(obja[0]				,							obja[1]						)
         self.assertIsNot(obja[0]				,							_lowerCAmelCase						)
         self.assertIsNot(obja[0]				,							obja[1]						)
         __snake_case:		str					    =		Hasher.hash(_lowerCAmelCase						)
         __snake_case:		Dict					    =		Hasher.hash(_lowerCAmelCase						)
         __snake_case:		Optional[Any]					    =		Hasher.hash(_lowerCAmelCase						)
         self.assertEqual(_lowerCAmelCase				,							_lowerCAmelCase						)
         self.assertEqual(_lowerCAmelCase				,							_lowerCAmelCase						)




       @require_tiktoken
       def 				UpperCAmelCase__					(    self					:       Any						):
         import tiktoken

         __snake_case:		Any					    =		tiktoken.get_encoding("""gpt2"""						)
         __snake_case:		List[Any]					    =		mda(datasets.utils.py_utils.dumps(_lowerCAmelCase						)						).hexdigest()
         __snake_case:		Optional[Any]					    =		tiktoken.get_encoding("""r50k_base"""						)
         __snake_case:		Optional[Any]					    =		mda(datasets.utils.py_utils.dumps(_lowerCAmelCase						)						).hexdigest()
         __snake_case:		Dict					    =		tiktoken.get_encoding("""gpt2"""						)
         __snake_case:		Dict					    =		mda(datasets.utils.py_utils.dumps(_lowerCAmelCase						)						).hexdigest()
         self.assertEqual(_lowerCAmelCase				,							_lowerCAmelCase						)
         self.assertNotEqual(_lowerCAmelCase				,							_lowerCAmelCase						)




       @require_torch
       def 				UpperCAmelCase__					(    self					:       Optional[Any]						):
         import torch

         __snake_case:		str					    =		torch.tensor([1.0]						)
         __snake_case:		Tuple					    =		mda(datasets.utils.py_utils.dumps(_lowerCAmelCase						)						).hexdigest()
         __snake_case:		int					    =		torch.tensor([2.0]						)
         __snake_case:		Any					    =		mda(datasets.utils.py_utils.dumps(_lowerCAmelCase						)						).hexdigest()
         __snake_case:		Any					    =		torch.tensor([1.0]						)
         __snake_case:		str					    =		mda(datasets.utils.py_utils.dumps(_lowerCAmelCase						)						).hexdigest()
         self.assertEqual(_lowerCAmelCase				,							_lowerCAmelCase						)
         self.assertNotEqual(_lowerCAmelCase				,							_lowerCAmelCase						)






       @require_spacy
       @require_spacy_model("""en_core_web_sm"""						)
       @require_spacy_model("""fr_core_news_sm"""						)
       @pytest.mark.integration
       def 				UpperCAmelCase__					(    self					:       Optional[int]						):
         import spacy

         __snake_case:		List[Any]					    =		spacy.load("""en_core_web_sm"""						)
         __snake_case:		List[Any]					    =		mda(datasets.utils.py_utils.dumps(_lowerCAmelCase						)						).hexdigest()
         __snake_case:		Dict					    =		spacy.load("""fr_core_news_sm"""						)
         __snake_case:		Optional[Any]					    =		mda(datasets.utils.py_utils.dumps(_lowerCAmelCase						)						).hexdigest()
         __snake_case:		Union[str, Any]					    =		spacy.load("""en_core_web_sm"""						)
         __snake_case:		Dict					    =		mda(datasets.utils.py_utils.dumps(_lowerCAmelCase						)						).hexdigest()
         self.assertEqual(_lowerCAmelCase				,							_lowerCAmelCase						)
         self.assertNotEqual(_lowerCAmelCase				,							_lowerCAmelCase						)







@pytest.mark.integration
def      A__				(       SCREAMING_SNAKE_CASE__)  ->       Union[str, Any]:
  __snake_case:		List[Any]					    =		tmp_path / "dir1"
  __snake_case:		Dict					    =		tmp_path / "dir2"
  dira.mkdir()
  dira.mkdir()
  __snake_case:		Optional[int]					    =		"script.py"
  __snake_case:		Union[str, Any]					    =		dedent(
      """\n    from datasets.fingerprint import Hasher\n    def foo():\n        pass\n    print(Hasher.hash(foo))\n    """)
  __snake_case:		Optional[int]					    =		dira / script_filename
  __snake_case:		str					    =		dira / script_filename
  with script_patha.open("""w""") as f:
    f.write(SCREAMING_SNAKE_CASE__)
  with script_patha.open("""w""") as f:
    f.write(SCREAMING_SNAKE_CASE__)
  __snake_case:		int					    =		subprocess.check_output(["""python""", str(SCREAMING_SNAKE_CASE__)])
  __snake_case:		int					    =		subprocess.check_output(["""python""", str(SCREAMING_SNAKE_CASE__)])
  assert fingerprinta == fingerprinta







def      A__				(       )  ->       int:
  __snake_case:		Union[str, Any]					    =		{"a": [0, 1, 2]}
  __snake_case:		Optional[int]					    =		DatasetChild(InMemoryTable.from_pydict(SCREAMING_SNAKE_CASE__))
  __snake_case:		Union[str, Any]					    =		dataset.funca()._fingerprint
  assert expected_fingerprint == dataset.funca()._fingerprint
  assert expected_fingerprint != dataset.funca()._fingerprint

  with Pool(2) as p:
    assert expected_fingerprint == p.apply_async(dataset.funca).get()._fingerprint
    assert expected_fingerprint != p.apply_async(dataset.funca).get()._fingerprint







def      A__				(       )  ->       List[str]:
  __snake_case:		Optional[Any]					    =		{"a": [0, 1, 2]}







  class        __snake_case			(				datasets.Dataset     ):





         '''simple docstring'''







         @fingerprint_transform(inplace=_lowerCAmelCase						)
         def 				UpperCAmelCase__					(    self					:       Any				,							A					:       Tuple						):
           return DummyDatasetChild(self.data				,							fingerprint=_lowerCAmelCase						)






  __snake_case:		Optional[int]					    =		DummyDatasetChild(InMemoryTable.from_pydict(SCREAMING_SNAKE_CASE__)).func()







  class        __snake_case			(				datasets.Dataset     ):





         '''simple docstring'''







         @fingerprint_transform(inplace=_lowerCAmelCase				,							version="""1.0.0"""						)
         def 				UpperCAmelCase__					(    self					:       List[str]				,							A					:       List[str]						):
           return DummyDatasetChild(self.data				,							fingerprint=_lowerCAmelCase						)






  __snake_case:		str					    =		DummyDatasetChild(InMemoryTable.from_pydict(SCREAMING_SNAKE_CASE__)).func()







  class        __snake_case			(				datasets.Dataset     ):





         '''simple docstring'''







         @fingerprint_transform(inplace=_lowerCAmelCase				,							version="""2.0.0"""						)
         def 				UpperCAmelCase__					(    self					:       Tuple				,							A					:       Union[str, Any]						):
           return DummyDatasetChild(self.data				,							fingerprint=_lowerCAmelCase						)






  __snake_case:		Optional[Any]					    =		DummyDatasetChild(InMemoryTable.from_pydict(SCREAMING_SNAKE_CASE__)).func()

  assert len({fingeprint_no_version, fingeprint_a, fingeprint_a}) == 3







def      A__				(       )  ->       Any:
  # AttributeError: module 'dill._dill' has no attribute 'stack'
  __snake_case:		Optional[int]					    =		Hasher()
  hasher.update(lambda SCREAMING_SNAKE_CASE__: x)












import logging
from pathlib import Path

import numpy as np
import pytorch_lightning as pl
import torch
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from pytorch_lightning.utilities import rank_zero_only

from utils import save_json



def    UpperCAmelCase__	(  lowerCamelCase ):
  lowercase     :Optional[Any]   						=    filter(lambda lowerCamelCase : p.requires_grad,      model.parameters() )
  lowercase     :Optional[int]   						=    sum([np.prod(p.size() ) for p in model_parameters] )
  return params


_UpperCAmelCase						:    Optional[Any]      =  logging.getLogger(__name__)




class  __lowerCAmelCase							(				pl.Callback):






       def        SCREAMING_SNAKE_CASE     ( self:       Union[str, Any]							,       _lowerCAmelCase:       List[Any]							,       _lowerCAmelCase:       Union[str, Any]					):
         lowercase     :Dict   						=    {F"lr_group_{i}": param["lr"] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups					)}
         pl_module.logger.log_metrics(_lowerCAmelCase					)






       @rank_zero_only
       def        SCREAMING_SNAKE_CASE     ( self:       Optional[Any]							,       _lowerCAmelCase:       pl.Trainer							,       _lowerCAmelCase:       pl.LightningModule							,       _lowerCAmelCase:       str							,       _lowerCAmelCase:       List[str]=True					):
         logger.info(F"***** {type_path} results at step {trainer.global_step:05d} *****"					)
         lowercase     :Tuple   						=    trainer.callback_metrics
         trainer.logger.log_metrics({k: v for k, v in metrics.items() if k not in ["log", "progress_bar", "preds"]}					)
         # Log results
         lowercase     :Optional[int]   						=    Path(pl_module.hparams.output_dir					)
         if type_path == "test":
           lowercase     :List[Any]   						=    od / "test_results.txt"
           lowercase     :Optional[int]   						=    od / "test_generations.txt"
         else:
           # this never gets hit. I prefer not to save intermediate generations, and results are in metrics.json
           # If people want this it will be easy enough to add back.
           lowercase     :int   						=    od / F"{type_path}_results/{trainer.global_step:05d}.txt"
           lowercase     :Tuple   						=    od / F"{type_path}_generations/{trainer.global_step:05d}.txt"
           results_file.parent.mkdir(exist_ok=_lowerCAmelCase					)
           generations_file.parent.mkdir(exist_ok=_lowerCAmelCase					)
         with open(_lowerCAmelCase							,       "a+"					) as writer:
           for key in sorted(_lowerCAmelCase					):
             if key in ["log", "progress_bar", "preds"]:
               continue
             lowercase     :int   						=    metrics[key]
             if isinstance(_lowerCAmelCase							,       torch.Tensor					):
               lowercase     :int   						=    val.item()
             lowercase     :Optional[Any]   						=    F"{key}: {val:.6f}\n"
             writer.write(_lowerCAmelCase					)

         if not save_generations:
           return

         if "preds" in metrics:
           lowercase     :str   						=    "\n".join(metrics["preds"]					)
           generations_file.open("w+"					).write(_lowerCAmelCase					)






       @rank_zero_only
       def        SCREAMING_SNAKE_CASE     ( self:       Optional[int]							,       _lowerCAmelCase:       Tuple							,       _lowerCAmelCase:       Tuple					):
         try:
           lowercase     :Optional[int]   						=    pl_module.model.model.num_parameters()
         except AttributeError:
           lowercase     :Dict   						=    pl_module.model.num_parameters()

         lowercase     :Any   						=    count_trainable_parameters(_lowerCAmelCase					)
         # mp stands for million parameters
         trainer.logger.log_metrics({"n_params": npars, "mp": npars / 1e6, "grad_mp": n_trainable_pars / 1e6}					)






       @rank_zero_only
       def        SCREAMING_SNAKE_CASE     ( self:       Union[str, Any]							,       _lowerCAmelCase:       pl.Trainer							,       _lowerCAmelCase:       pl.LightningModule					):
         save_json(pl_module.metrics							,       pl_module.metrics_save_path					)
         return self._write_logs(_lowerCAmelCase							,       _lowerCAmelCase							,       "test"					)






       @rank_zero_only
       def        SCREAMING_SNAKE_CASE     ( self:       List[str]							,       _lowerCAmelCase:       pl.Trainer							,       _lowerCAmelCase:       Dict					):
         save_json(pl_module.metrics							,       pl_module.metrics_save_path					)
         # Uncommenting this will save val generations
         # return self._write_logs(trainer, pl_module, "valid")







def    UpperCAmelCase__	(  lowerCamelCase,      lowerCamelCase,      lowerCamelCase=1,      lowerCamelCase=False ):
  if metric == "rouge2":
    lowercase     :str   						=    "{val_avg_rouge2:.4f}-{step_count}"
  elif metric == "bleu":
    lowercase     :int   						=    "{val_avg_bleu:.4f}-{step_count}"
  elif metric == "loss":
    lowercase     :List[str]   						=    "{val_avg_loss:.4f}-{step_count}"
  else:
    raise NotImplementedError(
        F"seq2seq callbacks only support rouge2, bleu and loss, got {metric}, You can make your own by adding to"
        " this function." )

  lowercase     :int   						=    ModelCheckpoint(
      dirpath=lowerCamelCase,      filename=lowerCamelCase,      monitor=F"val_{metric}",      mode="min" if "loss" in metric else "max",      save_top_k=lowerCamelCase,      )
  return checkpoint_callback







def    UpperCAmelCase__	(  lowerCamelCase,      lowerCamelCase ):
  return EarlyStopping(
      monitor=F"val_{metric}",      mode="min" if "loss" in metric else "max",      patience=lowerCamelCase,      verbose=lowerCamelCase,      )



