import csv
import json
import logging
import os
import re
import subprocess
import tempfile
import unittest
import zipfile
from pathlib import Path
from typing import Optional
from unittest import mock

import torch

# We use TF to parse the logs
from accelerate import Accelerator
from accelerate.test_utils.testing import (
    MockingTestCase,
    TempDirTestCase,
    require_comet_ml,
    require_tensorboard,
    require_wandb,
    skip,
)
from accelerate.tracking import CometMLTracker, GeneralTracker
from accelerate.utils import ProjectConfiguration, is_comet_ml_available, is_tensorboard_available


if is_comet_ml_available():
  from comet_ml import OfflineExperiment

if is_tensorboard_available():
  import struct

  import tensorboard.compat.proto.event_pba as event_pba

UpperCAmelCase				 =   logging.getLogger(__name__)






@require_tensorboard
class 				_SCREAMING_SNAKE_CASE( unittest.TestCase						):
   def _lowercase(     self     )						->Any:
        a	:     List[Any]						      =    '''test_project_with_config'''
        with tempfile.TemporaryDirectory() as dirpath:
             a	:     Union[str, Any]						      =    Accelerator(log_with='''tensorboard'''				,  project_dir=lowerCAmelCase_     )
             a	:     int						      =    {'''num_iterations''': 12, '''learning_rate''': 1e-2, '''some_boolean''': False, '''some_string''': '''some_value'''}
             accelerator.init_trackers(lowerCAmelCase_				,  lowerCAmelCase_     )
             accelerator.end_training()
             for child in Path(F"""{dirpath}/{project_name}"""     ).glob('''*/**'''     ):
                  a	:     Optional[int]						      =    list(filter(lambda lowerCAmelCase_     : x.is_file()				,  child.iterdir()     )     )[0]
             self.assertNotEqual(str(lowerCAmelCase_     )				,  ''''''     )
   def _lowercase(     self     )						->Optional[int]:
        a	:     str						      =    '''test_project_with_log'''
        with tempfile.TemporaryDirectory() as dirpath:
             a	:     List[str]						      =    Accelerator(log_with='''tensorboard'''				,  project_dir=lowerCAmelCase_     )
             accelerator.init_trackers(lowerCAmelCase_     )
             a	:     Union[str, Any]						      =    {'''total_loss''': 0.1, '''iteration''': 1, '''my_text''': '''some_value'''}
             accelerator.log(lowerCAmelCase_				,  step=0     )
             accelerator.end_training()
             # Logged values are stored in the outermost-tfevents file and can be read in as a TFRecord
             # Names are randomly generated each time
             a	:     List[Any]						      =    list(filter(lambda lowerCAmelCase_     : x.is_file()				,  Path(F"""{dirpath}/{project_name}"""     ).iterdir()     )     )[0]
             self.assertNotEqual(str(lowerCAmelCase_     )				,  ''''''     )
   def _lowercase(     self     )						->Union[str, Any]:
        a	:     Dict						      =    '''test_project_with_log'''
        with tempfile.TemporaryDirectory() as dirpath:
             a	:     List[str]						      =    Accelerator(log_with='''tensorboard'''				,  project_dir=lowerCAmelCase_     )
             accelerator.init_trackers(lowerCAmelCase_     )
             a	:     int						      =    {'''tensor''': torch.tensor(1     )}
             accelerator.log(lowerCAmelCase_				,  step=0     )
             accelerator.end_training()
             # Logged values are stored in the outermost-tfevents file and can be read in as a TFRecord
             # Names are randomly generated each time
             a	:     int						      =    list(filter(lambda lowerCAmelCase_     : x.is_file()				,  Path(F"""{dirpath}/{project_name}"""     ).iterdir()     )     )[0]
             # Reading implementation based on https://github.com/pytorch/pytorch/issues/45327#issuecomment-703757685
             with open(lowerCAmelCase_				,  '''rb'''     ) as f:
                  a	:     Tuple						      =    f.read()
             a	:     Dict						      =    False
             while data:
                  a	:     List[Any]						      =    struct.unpack('''Q'''				,  data[:8]     )

                  a	:     List[str]						      =    data[12 : 12 + int(header[0]     )]  # 8+4
                  a	:     Tuple						      =    data[12 + int(header[0]     ) + 4 :]
                  a	:     Dict						      =    event_pba.Event()

                  event.ParseFromString(lowerCAmelCase_     )
                  if event.HasField('''summary'''     ):
                       for value in event.summary.value:
                            if value.simple_value == 1.0 and value.tag == "tensor":
                                 a	:     str						      =    True
             self.assertTrue(lowerCAmelCase_				,  '''Converted tensor was not found in the log file!'''     )
   def _lowercase(     self     )						->Optional[int]:
        with self.assertRaisesRegex(lowerCAmelCase_				,  '''Logging with `tensorboard` requires a `logging_dir`'''     ):
             a	:     List[Any]						      =    Accelerator(log_with='''tensorboard'''     )
        with tempfile.TemporaryDirectory() as dirpath:
             a	:     int						      =    Accelerator(log_with='''tensorboard'''				,  project_dir=lowerCAmelCase_     )


   def _lowercase(     self     )						->List[Any]:
        a	:     Dict						      =    ProjectConfiguration(total_limit=30     )
        with tempfile.TemporaryDirectory() as dirpath:
             a	:     List[str]						      =    Accelerator(log_with='''tensorboard'''				,  project_dir=lowerCAmelCase_				,  project_config=lowerCAmelCase_     )






@require_wandb
@mock.patch.dict(os.environ							,   {'WANDB_MODE': 'offline'}						)
class 				_SCREAMING_SNAKE_CASE( snake_case__							,   snake_case__						):
   def _lowercase(     self     )						->Union[str, Any]:
        super().setUp()
        # wandb let's us override where logs are stored to via the WANDB_DIR env var
        self.add_mocks(mock.patch.dict(os.environ				,  {'''WANDB_DIR''': self.tmpdir}     )     )
   @staticmethod
   def _lowercase(     lowerCAmelCase_				,  lowerCAmelCase_				,  lowerCAmelCase_ = True     )						->str:
        a	:     Dict						      =    rF"""{section} ([\S\s]*?)\n\n"""
        if record:
             a	:     str						      =    rF"""Record: {pattern}"""
        a	:     Optional[int]						      =    re.findall(lowerCAmelCase_				,  lowerCAmelCase_     )[0]
        # A config
        if section == "config" or section == "history":
             a	:     Optional[int]						      =    re.findall(r'''"([a-zA-Z0-9_.,]+)'''				,  lowerCAmelCase_     )
             return {key: val for key, val in zip(cleaned_record[0::2]				,  cleaned_record[1::2]     )}
        # Everything else
        else:
             return dict(re.findall(r'''(\w+): "([^\s]+)"'''				,  lowerCAmelCase_     )     )


   @skip
   def _lowercase(     self     )						->Tuple:
        a	:     Any						      =    '''test_project_with_config'''
        a	:     Union[str, Any]						      =    Accelerator(log_with='''wandb'''     )
        a	:     List[str]						      =    {'''num_iterations''': 12, '''learning_rate''': 1e-2, '''some_boolean''': False, '''some_string''': '''some_value'''}
        a	:     Dict						      =    {'''wandb''': {'''tags''': ['''my_tag''']}}
        accelerator.init_trackers(lowerCAmelCase_				,  lowerCAmelCase_				,  lowerCAmelCase_     )
        a	:     int						      =    {'''total_loss''': 0.1, '''iteration''': 1, '''my_text''': '''some_value'''}
        accelerator.log(lowerCAmelCase_				,  step=0     )
        accelerator.end_training()
        # The latest offline log is stored at wandb/latest-run/*.wandb
        for child in Path(F"""{self.tmpdir}/wandb/latest-run"""     ).glob('''*'''     ):
             if child.is_file() and child.suffix == ".wandb":
                  a	:     Union[str, Any]						      =    subprocess.check_output(
                      ['''wandb''', '''sync''', '''--view''', '''--verbose''', str(lowerCAmelCase_     )]				,  env=os.environ.copy()     ).decode('''utf8'''				,  '''ignore'''     )
                  break

        # Check HPS through careful parsing and cleaning
        a	:     List[str]						      =    self.parse_log(lowerCAmelCase_				,  '''config'''     )
        self.assertEqual(logged_items['''num_iterations''']				,  '''12'''     )
        self.assertEqual(logged_items['''learning_rate''']				,  '''0.01'''     )
        self.assertEqual(logged_items['''some_boolean''']				,  '''false'''     )
        self.assertEqual(logged_items['''some_string''']				,  '''some_value'''     )
        self.assertEqual(logged_items['''some_string''']				,  '''some_value'''     )

        # Run tags
        a	:     Union[str, Any]						      =    self.parse_log(lowerCAmelCase_				,  '''run'''				,  lowerCAmelCase_     )
        self.assertEqual(logged_items['''tags''']				,  '''my_tag'''     )

        # Actual logging
        a	:     Dict						      =    self.parse_log(lowerCAmelCase_				,  '''history'''     )
        self.assertEqual(logged_items['''total_loss''']				,  '''0.1'''     )
        self.assertEqual(logged_items['''iteration''']				,  '''1'''     )
        self.assertEqual(logged_items['''my_text''']				,  '''some_value'''     )
        self.assertEqual(logged_items['''_step''']				,  '''0'''     )







def 						UpperCAmelCase							(		self    ,   lowerCAmelCase_    ,   lowerCAmelCase_							)						->						str:
     a	:     Dict						      =    run_name
     a	:     Dict						      =    OfflineExperiment(project_name=lowerCAmelCase_    ,   offline_directory=lowerCAmelCase_							)
     logger.info(f"""Initialized offline CometML project {self.run_name}"""							)
     logger.info('''Make sure to log any initial configurations with `self.store_init_configuration` before training!'''							)






@require_comet_ml
@mock.patch.object(snake_case__							,   '__init__'							,   snake_case__						)
class 				_SCREAMING_SNAKE_CASE( unittest.TestCase						):
   @staticmethod
   def _lowercase(     lowerCAmelCase_				,  lowerCAmelCase_				,  lowerCAmelCase_ = False     )						->Tuple:
        for log in log_list:
             a	:     str						      =    json.loads(lowerCAmelCase_     )['''payload''']
             if is_param and "param" in j.keys():
                  if j["param"]["paramName"] == key:
                       return j["param"]["paramValue"]
             if "log_other" in j.keys():
                  if j["log_other"]["key"] == key:
                       return j["log_other"]["val"]
             if "metric" in j.keys():
                  if j["metric"]["metricName"] == key:
                       return j["metric"]["metricValue"]
   def _lowercase(     self     )						->Dict:
        with tempfile.TemporaryDirectory() as d:
             a	:     List[str]						      =    CometMLTracker('''test_project_with_config'''				,  lowerCAmelCase_     )
             a	:     List[str]						      =    Accelerator(log_with=lowerCAmelCase_     )
             a	:     int						      =    {'''num_iterations''': 12, '''learning_rate''': 1e-2, '''some_boolean''': False, '''some_string''': '''some_value'''}
             accelerator.init_trackers(lowerCAmelCase_				,  lowerCAmelCase_     )
             accelerator.end_training()
             a	:     int						      =    os.listdir(lowerCAmelCase_     )[0]  # Comet is nice, it's just a zip file here
             # We parse the raw logs
             a	:     Optional[Any]						      =    os.path.join(lowerCAmelCase_				,  lowerCAmelCase_     )
             a	:     List[Any]						      =    zipfile.ZipFile(lowerCAmelCase_				,  '''r'''     )
             a	:     Union[str, Any]						      =    archive.open('''messages.json'''     ).read().decode('''utf-8'''     )
        a	:     Optional[Any]						      =    log.split('''\n'''     )[:-1]
        self.assertEqual(self.get_value_from_key(lowerCAmelCase_				,  '''num_iterations'''				,  lowerCAmelCase_     )				,  12     )
        self.assertEqual(self.get_value_from_key(lowerCAmelCase_				,  '''learning_rate'''				,  lowerCAmelCase_     )				,  0.01     )
        self.assertEqual(self.get_value_from_key(lowerCAmelCase_				,  '''some_boolean'''				,  lowerCAmelCase_     )				,  lowerCAmelCase_     )
        self.assertEqual(self.get_value_from_key(lowerCAmelCase_				,  '''some_string'''				,  lowerCAmelCase_     )				,  '''some_value'''     )


   def _lowercase(     self     )						->Tuple:
        with tempfile.TemporaryDirectory() as d:
             a	:     List[str]						      =    CometMLTracker('''test_project_with_config'''				,  lowerCAmelCase_     )
             a	:     Dict						      =    Accelerator(log_with=lowerCAmelCase_     )
             accelerator.init_trackers(lowerCAmelCase_     )
             a	:     Any						      =    {'''total_loss''': 0.1, '''iteration''': 1, '''my_text''': '''some_value'''}
             accelerator.log(lowerCAmelCase_				,  step=0     )
             accelerator.end_training()
             a	:     Any						      =    os.listdir(lowerCAmelCase_     )[0]  # Comet is nice, it's just a zip file here
             # We parse the raw logs
             a	:     Tuple						      =    os.path.join(lowerCAmelCase_				,  lowerCAmelCase_     )
             a	:     Tuple						      =    zipfile.ZipFile(lowerCAmelCase_				,  '''r'''     )
             a	:     List[str]						      =    archive.open('''messages.json'''     ).read().decode('''utf-8'''     )
        a	:     Any						      =    log.split('''\n'''     )[:-1]
        self.assertEqual(self.get_value_from_key(lowerCAmelCase_				,  '''curr_step'''				,  lowerCAmelCase_     )				,  0     )
        self.assertEqual(self.get_value_from_key(lowerCAmelCase_				,  '''total_loss'''     )				,  0.1     )
        self.assertEqual(self.get_value_from_key(lowerCAmelCase_				,  '''iteration'''     )				,  1     )
        self.assertEqual(self.get_value_from_key(lowerCAmelCase_				,  '''my_text'''     )				,  '''some_value'''     )






class 				_SCREAMING_SNAKE_CASE( snake_case__						):
   UpperCamelCase_					:    Optional[Any]			      = [
       'total_loss',
       'iteration',
       'my_text',
       'learning_rate',
       'num_iterations',
       'some_boolean',
       'some_string',
   ]

   UpperCamelCase_					:    List[Any]			      = 'my_custom_tracker'
   UpperCamelCase_					:    Optional[int]			      = False
   def __init__(     self				,  lowerCAmelCase_     )						->Tuple:
        a	:     Optional[Any]						      =    open(F"""{dir}/log.csv"""				,  '''w+'''     )
        a	:     Optional[Any]						      =    csv.DictWriter(self.f				,  fieldnames=self._col_names     )
        self.writer.writeheader()
   @property
   def _lowercase(     self     )						->Tuple:
        return self.writer
   def _lowercase(     self				,  lowerCAmelCase_     )						->List[str]:
        logger.info('''Call init'''     )
        self.writer.writerow(lowerCAmelCase_     )
   def _lowercase(     self				,  lowerCAmelCase_				,  lowerCAmelCase_     )						->List[str]:
        logger.info('''Call log'''     )
        self.writer.writerow(lowerCAmelCase_     )


   def _lowercase(     self     )						->int:
        self.f.close()






class 				_SCREAMING_SNAKE_CASE( unittest.TestCase						):
   def _lowercase(     self     )						->Tuple:
        with tempfile.TemporaryDirectory() as d:
             a	:     Optional[Any]						      =    MyCustomTracker(lowerCAmelCase_     )
             a	:     Any						      =    Accelerator(log_with=lowerCAmelCase_     )
             a	:     Any						      =    {'''num_iterations''': 12, '''learning_rate''': 1e-2, '''some_boolean''': False, '''some_string''': '''some_value'''}
             accelerator.init_trackers('''Some name'''				,  lowerCAmelCase_     )
             accelerator.end_training()
             with open(F"""{d}/log.csv"""				,  '''r'''     ) as f:
                  a	:     Any						      =    csv.DictReader(lowerCAmelCase_     )
                  a	:     Optional[int]						      =    next(lowerCAmelCase_     )
                  a	:     List[Any]						      =    {
                      '''total_loss''': '''''',
                      '''iteration''': '''''',
                      '''my_text''': '''''',
                      '''learning_rate''': '''0.01''',
                      '''num_iterations''': '''12''',
                      '''some_boolean''': '''False''',
                      '''some_string''': '''some_value''',
                  }
                  self.assertDictEqual(lowerCAmelCase_				,  lowerCAmelCase_     )


   def _lowercase(     self     )						->Optional[int]:
        with tempfile.TemporaryDirectory() as d:
             a	:     str						      =    MyCustomTracker(lowerCAmelCase_     )
             a	:     str						      =    Accelerator(log_with=lowerCAmelCase_     )
             accelerator.init_trackers('''Some name'''     )
             a	:     Any						      =    {'''total_loss''': 0.1, '''iteration''': 1, '''my_text''': '''some_value'''}
             accelerator.log(lowerCAmelCase_				,  step=0     )
             accelerator.end_training()
             with open(F"""{d}/log.csv"""				,  '''r'''     ) as f:
                  a	:     Dict						      =    csv.DictReader(lowerCAmelCase_     )
                  a	:     int						      =    next(lowerCAmelCase_     )
                  a	:     Dict						      =    {
                      '''total_loss''': '''0.1''',
                      '''iteration''': '''1''',
                      '''my_text''': '''some_value''',
                      '''learning_rate''': '''''',
                      '''num_iterations''': '''''',
                      '''some_boolean''': '''''',
                      '''some_string''': '''''',
                  }
                  self.assertDictEqual(lowerCAmelCase_				,  lowerCAmelCase_     )
import re
import tempfile
import unittest
from pathlib import Path

from datasets.utils.metadata import DatasetMetadata







def 						UpperCAmelCase							(		lowerCAmelCase_							)						->						str:
     a	:     Optional[Any]						      =    min(re.search('''^ +'''    ,   lowerCAmelCase_							).end() if t.startswith(''' '''							) else 0 for t in string.splitlines()							)
     return "\n".join([line[indent_level:] for line in string.splitlines() if indent_level < len(lowerCAmelCase_							)]							)


UpperCAmelCase				 =   """\
---
language:
- zh
- en
task_ids:
- sentiment-classification
---
# Begin of markdown

Some cool dataset card
"""

UpperCAmelCase				 =   """\
---
---
# Begin of markdown

Some cool dataset card
"""


UpperCAmelCase				 =   """\
# Begin of markdown

Some cool dataset card
"""






class 				_SCREAMING_SNAKE_CASE( unittest.TestCase						):
   def _lowercase(     self     )						->List[str]:
        with tempfile.TemporaryDirectory() as tmp_dir:
             a	:     List[str]						      =    Path(lowerCAmelCase_     ) / '''README.md'''
             with open(lowerCAmelCase_				,  '''w+'''     ) as readme_file:
                  readme_file.write(lowerCAmelCase_     )
             a	:     Dict						      =    DatasetMetadata.from_readme(lowerCAmelCase_     )
             self.assertDictEqual(lowerCAmelCase_				,  {'''language''': ['''zh''', '''en'''], '''task_ids''': ['''sentiment-classification''']}     )

             with open(lowerCAmelCase_				,  '''w+'''     ) as readme_file:
                  readme_file.write(lowerCAmelCase_     )
             a	:     List[str]						      =    DatasetMetadata.from_readme(lowerCAmelCase_     )
             self.assertDictEqual(lowerCAmelCase_				,  {}     )

             with open(lowerCAmelCase_				,  '''w+'''     ) as readme_file:
                  readme_file.write(lowerCAmelCase_     )
             a	:     List[Any]						      =    DatasetMetadata.from_readme(lowerCAmelCase_     )
             self.assertEqual(lowerCAmelCase_				,  {}     )


   def _lowercase(     self     )						->Optional[Any]:
        a	:     Tuple						      =    _dedent(
            '''\
            annotations_creators:
            - found
            language_creators:
            - found
            language:
            - en
            license:
            - unknown
            multilinguality:
            - monolingual
            pretty_name: Test Dataset
            size_categories:
            - 10K<n<100K
            source_datasets:
            - extended|other-yahoo-webscope-l6
            task_categories:
            - question-answering
            task_ids:
            - open-domain-qa
            '''     )
        assert DatasetMetadata.from_yaml_string(lowerCAmelCase_     )

        a	:     Dict						      =    _dedent(
            '''\
            annotations_creators:
            - found
            language:
            - en
            license:
            - unknown
            multilinguality:
            - monolingual
            pretty_name: Test Dataset
            size_categories:
            - 10K<n<100K
            source_datasets:
            - extended|other-yahoo-webscope-l6
            task_categories:
            - question-answering
            task_ids:
            - open-domain-qa
            task_ids:
            - open-domain-qa
            '''     )
        with self.assertRaises(lowerCAmelCase_     ):
             DatasetMetadata.from_yaml_string(lowerCAmelCase_     )

        a	:     Optional[int]						      =    _dedent(
            '''\
            annotations_creators:
            - found
            language_creators:
            - found
            language:
            - en
            license:
            - unknown
            multilinguality:
            - monolingual
            pretty_name: Test Dataset
            size_categories:
            - 10K<n<100K
            source_datasets:
            - extended|other-yahoo-webscope-l6
            task_categories:
            - text-classification
            task_ids:
            - multi-class-classification
            paperswithcode_id:
            - squad
            configs:
            - en
            train-eval-index:
            - config: en
              task: text-classification
              task_id: multi_class_classification
              splits:
                train_split: train
                eval_split: test
              col_mapping:
                text: text
                label: target
              metrics:
                - type: accuracy
                  name: Accuracy
            extra_gated_prompt: |
              By clicking on “Access repository” below, you also agree to ImageNet Terms of Access:
              [RESEARCHER_FULLNAME] (the "Researcher") has requested permission to use the ImageNet database (the "Database") at Princeton University and Stanford University. In exchange for such permission, Researcher hereby agrees to the following terms and conditions:
              1. Researcher shall use the Database only for non-commercial research and educational purposes.
            extra_gated_fields:
              Company: text
              Country: text
              I agree to use this model for non-commerical use ONLY: checkbox
            '''     )
        assert DatasetMetadata.from_yaml_string(lowerCAmelCase_     )