





'''simple docstring'''




import itertools
import json
import linecache
import os
import pickle
import re
import socket
import string
from collections import Counter
from logging import getLogger
from pathlib import Path
from typing import Callable, Dict, Iterable, List

import git
import torch
from torch.utils.data import Dataset

from transformers import BartTokenizer, RagTokenizer, TaTokenizer







def _lowerCAmelCase     (      _lowerCAmelCase       ,    _lowerCAmelCase       ,    _lowerCAmelCase       ,    _lowerCAmelCase       ,    _lowerCAmelCase=True       ,    _lowerCAmelCase="pt"    )->     Optional[int]:
 __UpperCAmelCase					=   {'add_prefix_space': True} if isinstance(_lowerCAmelCase       ,    _lowerCAmelCase    ) and not line.startswith(' '    ) else {}
 __UpperCAmelCase					=   padding_side
 return tokenizer(
     [line]       ,    max_length=_lowerCAmelCase       ,    padding='max_length' if pad_to_max_length else None       ,    truncation=_lowerCAmelCase       ,    return_tensors=_lowerCAmelCase       ,    add_special_tokens=_lowerCAmelCase       ,    **_lowerCAmelCase       ,    )







def _lowerCAmelCase     (      _lowerCAmelCase       ,    _lowerCAmelCase       ,    _lowerCAmelCase=None       ,    )->     Union[str, Any]:
 __UpperCAmelCase					=   input_ids.ne(_lowerCAmelCase    ).any(dim=0    )
 if attention_mask is None:
  return input_ids[:, keep_column_mask]
 else:
  return (input_ids[:, keep_column_mask], attention_mask[:, keep_column_mask])






class UpperCAmelCase							(		UpperCAmelCase_							):



    def __init__(  self     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A="train"     ,			__A=None     ,			__A=None     ,			__A=None     ,			__A=""     ,			):
     super().__init__()
     __UpperCAmelCase					=   Path(__A					).joinpath(type_path + '.source'					)
     __UpperCAmelCase					=   Path(__A					).joinpath(type_path + '.target'					)
     __UpperCAmelCase					=   self.get_char_lens(self.src_file					)
     __UpperCAmelCase					=   max_source_length
     __UpperCAmelCase					=   max_target_length
     assert min(self.src_lens					) > 0, f'found empty line in {self.src_file}'
     __UpperCAmelCase					=   tokenizer
     __UpperCAmelCase					=   prefix
     if n_obs is not None:
      __UpperCAmelCase					=   self.src_lens[:n_obs]
     __UpperCAmelCase					=   src_lang
     __UpperCAmelCase					=   tgt_lang



    def __len__(  self					):
     return len(self.src_lens					)



    def __getitem__(  self     ,			__A					):
     __UpperCAmelCase					=   index + 1  # linecache starts at 1
     __UpperCAmelCase					=   self.prefix + linecache.getline(str(self.src_file					)     ,			__A					).rstrip('\n'					)
     __UpperCAmelCase					=   linecache.getline(str(self.tgt_file					)     ,			__A					).rstrip('\n'					)
     assert source_line, f'empty source line for index {index}'
     assert tgt_line, f'empty tgt line for index {index}'

     # Need to add eos token manually for T5
     if isinstance(self.tokenizer     ,			__A					):
      source_line += self.tokenizer.eos_token
      tgt_line += self.tokenizer.eos_token

     # Pad source and target to the right
     __UpperCAmelCase					=   (
         self.tokenizer.question_encoder if isinstance(self.tokenizer     ,			__A					) else self.tokenizer
     )
     __UpperCAmelCase					=   self.tokenizer.generator if isinstance(self.tokenizer     ,			__A					) else self.tokenizer

     __UpperCAmelCase					=   encode_line(__A     ,			__A     ,			self.max_source_length     ,			'right'					)
     __UpperCAmelCase					=   encode_line(__A     ,			__A     ,			self.max_target_length     ,			'right'					)

     __UpperCAmelCase					=   source_inputs['input_ids'].squeeze()
     __UpperCAmelCase					=   target_inputs['input_ids'].squeeze()
     __UpperCAmelCase					=   source_inputs['attention_mask'].squeeze()
     return {
         "input_ids": source_ids,
         "attention_mask": src_mask,
         "decoder_input_ids": target_ids,
     }



    @staticmethod
    def 					__lowerCamelCase						(  __A					):
     return [len(__A					) for x in Path(__A					).open().readlines()]




    def 					__lowerCamelCase						(  self     ,			__A					):
     __UpperCAmelCase					=   torch.stack([x['input_ids'] for x in batch]					)
     __UpperCAmelCase					=   torch.stack([x['attention_mask'] for x in batch]					)
     __UpperCAmelCase					=   torch.stack([x['decoder_input_ids'] for x in batch]					)
     __UpperCAmelCase					=   (
         self.tokenizer.generator.pad_token_id
         if isinstance(self.tokenizer     ,			__A					)
         else self.tokenizer.pad_token_id
     )
     __UpperCAmelCase					=   (
         self.tokenizer.question_encoder.pad_token_id
         if isinstance(self.tokenizer     ,			__A					)
         else self.tokenizer.pad_token_id
     )
     __UpperCAmelCase					=   trim_batch(__A     ,			__A					)
     __UpperCAmelCase					,    __UpperCAmelCase					=   trim_batch(__A     ,			__A     ,			attention_mask=__A					)
     __UpperCAmelCase					=   {
         'input_ids': source_ids,
         'attention_mask': source_mask,
         'decoder_input_ids': y,
     }
     return batch



_A:					List[Any]							=				getLogger(__name__)







def _lowerCAmelCase     (      _lowerCAmelCase    )->     List[Any]:
 return list(itertools.chain.from_iterable(_lowerCAmelCase    )    )







def _lowerCAmelCase     (      _lowerCAmelCase    )->     None:
 __UpperCAmelCase					=   get_git_info()
 save_json(_lowerCAmelCase       ,    os.path.join(_lowerCAmelCase       ,    'git_log.json'    )    )







def _lowerCAmelCase     (      _lowerCAmelCase       ,    _lowerCAmelCase       ,    _lowerCAmelCase=4       ,    **_lowerCAmelCase    )->     List[str]:
 with open(_lowerCAmelCase       ,    'w'    ) as f:
  json.dump(_lowerCAmelCase       ,    _lowerCAmelCase       ,    indent=_lowerCAmelCase       ,    **_lowerCAmelCase    )







def _lowerCAmelCase     (      _lowerCAmelCase    )->     List[Any]:
 with open(_lowerCAmelCase    ) as f:
  return json.load(_lowerCAmelCase    )







def _lowerCAmelCase     (      )->     Union[str, Any]:
 __UpperCAmelCase					=   git.Repo(search_parent_directories=_lowerCAmelCase    )
 __UpperCAmelCase					=   {
     'repo_id': str(_lowerCAmelCase    ),
     'repo_sha': str(repo.head.object.hexsha    ),
     'repo_branch': str(repo.active_branch    ),
     'hostname': str(socket.gethostname()    ),
 }
 return repo_infos







def _lowerCAmelCase     (      _lowerCAmelCase       ,    _lowerCAmelCase    )->     List:
 return list(map(_lowerCAmelCase       ,    _lowerCAmelCase    )    )







def _lowerCAmelCase     (      _lowerCAmelCase       ,    _lowerCAmelCase    )->     Dict:
 with open(_lowerCAmelCase       ,    'wb'    ) as f:
  return pickle.dump(_lowerCAmelCase       ,    _lowerCAmelCase    )







def _lowerCAmelCase     (      _lowerCAmelCase    )->     int:

 def remove_articles(_lowerCAmelCase    ):
  return re.sub(r'\b(a|an|the)\b'       ,    ' '       ,    _lowerCAmelCase    )

 def white_space_fix(_lowerCAmelCase    ):
  return " ".join(text.split()    )

 def remove_punc(_lowerCAmelCase    ):
  __UpperCAmelCase					=   set(string.punctuation    )
  return "".join(ch for ch in text if ch not in exclude    )

 def lower(_lowerCAmelCase    ):
  return text.lower()

 return white_space_fix(remove_articles(remove_punc(lower(_lowerCAmelCase    )    )    )    )







def _lowerCAmelCase     (      _lowerCAmelCase       ,    _lowerCAmelCase    )->     int:
 __UpperCAmelCase					=   normalize_answer(_lowerCAmelCase    ).split()
 __UpperCAmelCase					=   normalize_answer(_lowerCAmelCase    ).split()
 __UpperCAmelCase					=   Counter(_lowerCAmelCase    ) & Counter(_lowerCAmelCase    )
 __UpperCAmelCase					=   sum(common.values()    )
 if num_same == 0:
  return 0
 __UpperCAmelCase					=   1.0 * num_same / len(_lowerCAmelCase    )
 __UpperCAmelCase					=   1.0 * num_same / len(_lowerCAmelCase    )
 __UpperCAmelCase					=   (2 * precision * recall) / (precision + recall)
 return fa







def _lowerCAmelCase     (      _lowerCAmelCase       ,    _lowerCAmelCase    )->     Union[str, Any]:
 return normalize_answer(_lowerCAmelCase    ) == normalize_answer(_lowerCAmelCase    )







def _lowerCAmelCase     (      _lowerCAmelCase       ,    _lowerCAmelCase    )->     Dict:
 assert len(_lowerCAmelCase    ) == len(_lowerCAmelCase    )
 __UpperCAmelCase					=   0
 for hypo, pred in zip(_lowerCAmelCase       ,    _lowerCAmelCase    ):
  em += exact_match_score(_lowerCAmelCase       ,    _lowerCAmelCase    )
 if len(_lowerCAmelCase    ) > 0:
  em /= len(_lowerCAmelCase    )
 return {"em": em}







def _lowerCAmelCase     (      _lowerCAmelCase    )->     Dict:
 return model_prefix.startswith('rag'    )







def _lowerCAmelCase     (      _lowerCAmelCase       ,    _lowerCAmelCase       ,    _lowerCAmelCase    )->     int:
 __UpperCAmelCase					=   {p: p for p in extra_params}
 # T5 models don't have `dropout` param, they have `dropout_rate` instead
 __UpperCAmelCase					=   'dropout_rate'
 for p in extra_params:
  if getattr(_lowerCAmelCase       ,    _lowerCAmelCase       ,    _lowerCAmelCase    ):
   if not hasattr(_lowerCAmelCase       ,    _lowerCAmelCase    ) and not hasattr(_lowerCAmelCase       ,    equivalent_param[p]    ):
    logger.info('config doesn\'t have a `{}` attribute'.format(_lowerCAmelCase    )    )
    delattr(_lowerCAmelCase       ,    _lowerCAmelCase    )
    continue
   __UpperCAmelCase					=   p if hasattr(_lowerCAmelCase       ,    _lowerCAmelCase    ) else equivalent_param[p]
   setattr(_lowerCAmelCase       ,    _lowerCAmelCase       ,    getattr(_lowerCAmelCase       ,    _lowerCAmelCase    )    )
   delattr(_lowerCAmelCase       ,    _lowerCAmelCase    )
 return hparams, config












'''simple docstring'''




import multiprocessing
from typing import TYPE_CHECKING, Optional, Union

from .. import Dataset, Features, config
from ..formatting import query_table
from ..packaged_modules.sql.sql import Sql
from ..utils import logging
from .abc import AbstractDatasetInputStream


if TYPE_CHECKING:
    import sqlitea

    import sqlalchemy






class UpperCAmelCase							(		UpperCAmelCase_							):



    def __init__(  self     ,			__A     ,			__A     ,			__A = None     ,			__A = None     ,			__A = False     ,			**__A     ,			):
     super().__init__(features=__A     ,			cache_dir=__A     ,			keep_in_memory=__A     ,			**__A					)
     __UpperCAmelCase					=   Sql(
         cache_dir=__A     ,			features=__A     ,			sql=__A     ,			con=__A     ,			**__A     ,			)




    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   None
     __UpperCAmelCase					=   None
     __UpperCAmelCase					=   None
     __UpperCAmelCase					=   None

     self.builder.download_and_prepare(
         download_config=__A     ,			download_mode=__A     ,			verification_mode=__A     ,			base_path=__A     ,			)

     # Build dataset for splits
     __UpperCAmelCase					=   self.builder.as_dataset(
         split='train'     ,			verification_mode=__A     ,			in_memory=self.keep_in_memory					)
     return dataset






class UpperCAmelCase							:



    def __init__(  self     ,			__A     ,			__A     ,			__A     ,			__A = None     ,			__A = None     ,			**__A     ,			):
     if num_proc is not None and num_proc <= 0:
      raise ValueError(f'num_proc {num_proc} must be an integer > 0.'					)

     __UpperCAmelCase					=   dataset
     __UpperCAmelCase					=   name
     __UpperCAmelCase					=   con
     __UpperCAmelCase					=   batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
     __UpperCAmelCase					=   num_proc
     __UpperCAmelCase					=   to_sql_kwargs



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.to_sql_kwargs.pop('sql'     ,			__A					)
     __UpperCAmelCase					=   self.to_sql_kwargs.pop('con'     ,			__A					)
     __UpperCAmelCase					=   self.to_sql_kwargs.pop('index'     ,			__A					)

     __UpperCAmelCase					=   self._write(index=__A     ,			**self.to_sql_kwargs					)
     return written



    def 					__lowerCamelCase						(  self     ,			__A					):
     __UpperCAmelCase					,    __UpperCAmelCase					,    __UpperCAmelCase					=   args
     __UpperCAmelCase					=   {**to_sql_kwargs, 'if_exists': 'append'} if offset > 0 else to_sql_kwargs
     __UpperCAmelCase					=   query_table(
         table=self.dataset.data     ,			key=slice(__A     ,			offset + self.batch_size					)     ,			indices=self.dataset._indices     ,			)
     __UpperCAmelCase					=   batch.to_pandas()
     __UpperCAmelCase					=   df.to_sql(self.name     ,			self.con     ,			index=__A     ,			**__A					)
     return num_rows or len(__A					)




    def 					__lowerCamelCase						(  self     ,			__A     ,			**__A					):
     __UpperCAmelCase					=   0

     if self.num_proc is None or self.num_proc == 1:
      for offset in logging.tqdm(
          range(0     ,			len(self.dataset					)     ,			self.batch_size					)     ,			unit='ba'     ,			disable=not logging.is_progress_bar_enabled()     ,			desc='Creating SQL from Arrow format'     ,			):
       written += self._batch_sql((offset, index, to_sql_kwargs)					)
     else:
      __UpperCAmelCase					,    __UpperCAmelCase					=   len(self.dataset					), self.batch_size
      with multiprocessing.Pool(self.num_proc					) as pool:
       for num_rows in logging.tqdm(
           pool.imap(
               self._batch_sql     ,			[(offset, index, to_sql_kwargs) for offset in range(0     ,			__A     ,			__A					)]     ,			)     ,			total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size     ,			unit='ba'     ,			disable=not logging.is_progress_bar_enabled()     ,			desc='Creating SQL from Arrow format'     ,			):
        written += num_rows

     return written





