




"""simple docstring"""

import argparse
import logging
import os
import sys
import time

import tensorflow as tf
from datasets import load_dataset
from tqdm import tqdm

from transformers import AutoTokenizer, TFAutoModelForSequenceClassification
from transformers.utils import is_sagemaker_dp_enabled


if os.environ.get("SDP_ENABLED") or is_sagemaker_dp_enabled():
      __lowerCamelCase								=						True
      __lowerCamelCase								=						"p3dn.24xlarge"
      import smdistributed.dataparallel.tensorflow as sdp
else:
      __lowerCamelCase								=						False



def 				a (						__snake_case     :				List[str],    __snake_case     :				Optional[Any],    __snake_case     :				Optional[int],    __snake_case     :				str,    __snake_case     :				int,    __snake_case     :				Any,    __snake_case     :				Any=None  ):



    '''simple docstring'''





    UpperCAmelCase_					:Tuple      			=      tqdm(__snake_case  )
    for i, batch in enumerate(__snake_case  ):
        with tf.GradientTape() as tape:
            UpperCAmelCase_					:List[str]      			=      batch
            UpperCAmelCase_					:Optional[int]      			=      model(__snake_case  )
            UpperCAmelCase_					:Union[str, Any]      			=      loss(__snake_case,    outputs.logits  )

        if SDP_ENABLED:
            UpperCAmelCase_					:str      			=      sdp.DistributedGradientTape(__snake_case,    sparse_as_dense=__snake_case  )

        UpperCAmelCase_					:Optional[int]      			=      tape.gradient(__snake_case,    model.trainable_variables  )
        opt.apply_gradients(zip(__snake_case,    model.trainable_variables  )  )

        pbar.set_description(f'Loss: {loss_value:.4f}'  )

        if SDP_ENABLED and i == 0:
            sdp.broadcast_variables(model.variables,    root_rank=0  )
            sdp.broadcast_variables(opt.variables(),    root_rank=0  )

        if max_steps and i >= max_steps:
            break

    UpperCAmelCase_					:str      			=      {'''loss''': loss_value.numpy()}
    return train_results



def 				a (						__snake_case     :				List[Any],    __snake_case     :				Optional[Any],    __snake_case     :				int  ):



    '''simple docstring'''





    UpperCAmelCase_					:Optional[int]      			=      load_dataset('''imdb''',    split=['''train''', '''test''']  )

    # Preprocess train dataset
    UpperCAmelCase_					:int      			=      train_dataset.map(
        lambda __snake_case  : tokenizer(e['''text'''],    truncation=__snake_case,    padding='''max_length'''  ),    batched=__snake_case  )
    train_dataset.set_format(type='''tensorflow''',    columns=['''input_ids''', '''attention_mask''', '''label''']  )

    UpperCAmelCase_					:Optional[int]      			=      {
        x: train_dataset[x].to_tensor(default_value=0,    shape=[None, tokenizer.model_max_length]  )
        for x in ['''input_ids''', '''attention_mask''']
    }
    UpperCAmelCase_					:str      			=      tf.data.Dataset.from_tensor_slices((train_features, train_dataset['''label'''])  )

    # Preprocess test dataset
    UpperCAmelCase_					:int      			=      test_dataset.map(
        lambda __snake_case  : tokenizer(e['''text'''],    truncation=__snake_case,    padding='''max_length'''  ),    batched=__snake_case  )
    test_dataset.set_format(type='''tensorflow''',    columns=['''input_ids''', '''attention_mask''', '''label''']  )

    UpperCAmelCase_					:Dict      			=      {
        x: test_dataset[x].to_tensor(default_value=0,    shape=[None, tokenizer.model_max_length]  )
        for x in ['''input_ids''', '''attention_mask''']
    }
    UpperCAmelCase_					:str      			=      tf.data.Dataset.from_tensor_slices((test_features, test_dataset['''label'''])  )

    if SDP_ENABLED:
        UpperCAmelCase_					:str      			=      tf_train_dataset.shard(sdp.size(),    sdp.rank()  )
        UpperCAmelCase_					:Any      			=      tf_test_dataset.shard(sdp.size(),    sdp.rank()  )
    UpperCAmelCase_					:List[Any]      			=      tf_train_dataset.batch(__snake_case,    drop_remainder=__snake_case  )
    UpperCAmelCase_					:Dict      			=      tf_test_dataset.batch(__snake_case,    drop_remainder=__snake_case  )

    return tf_train_dataset, tf_test_dataset


if __name__ == "__main__":
      __lowerCamelCase								=						argparse.ArgumentParser()

      # Hyperparameters sent by the client are passed as command-line arguments to the script.
      parser.add_argument("--epochs", type=int, default=3)
      parser.add_argument("--per_device_train_batch_size", type=int, default=16)
      parser.add_argument("--per_device_eval_batch_size", type=int, default=8)
      parser.add_argument("--model_name_or_path", type=str)
      parser.add_argument("--learning_rate", type=str, default=5E-5)
      parser.add_argument("--do_train", type=bool, default=True)
      parser.add_argument("--do_eval", type=bool, default=True)
      parser.add_argument("--output_dir", type=str)
      parser.add_argument("--max_steps", type=int, default=None)

      # Data, model, and output directories
      parser.add_argument("--output_data_dir", type=str, default=os.environ["SM_OUTPUT_DATA_DIR"])
      parser.add_argument("--model_dir", type=str, default=os.environ["SM_MODEL_DIR"])
      parser.add_argument("--n_gpus", type=str, default=os.environ["SM_NUM_GPUS"])

      __lowerCamelCase				,     __lowerCamelCase								=						parser.parse_known_args()

      # Set up logging
      __lowerCamelCase								=						logging.getLogger(__name__)

      logging.basicConfig(
          level=logging.getLevelName("INFO"),
          handlers=[logging.StreamHandler(sys.stdout)],
          format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
      )

      if SDP_ENABLED:
            sdp.init()

            __lowerCamelCase								=						tf.config.experimental.list_physical_devices("GPU")
            for gpu in gpus:
                  tf.config.experimental.set_memory_growth(gpu, True)
            if gpus:
                  tf.config.experimental.set_visible_devices(gpus[sdp.local_rank()], "GPU")

    # Load model and tokenizer
      __lowerCamelCase								=						TFAutoModelForSequenceClassification.from_pretrained(args.model_name_or_path)
      __lowerCamelCase								=						AutoTokenizer.from_pretrained(args.model_name_or_path)

      # get datasets
      __lowerCamelCase				,     __lowerCamelCase								=						get_datasets(
          tokenizer=tokenizer,
          train_batch_size=args.per_device_train_batch_size,
          eval_batch_size=args.per_device_eval_batch_size,
      )

      # fine optimizer and loss
      __lowerCamelCase								=						tf.keras.optimizers.Adam(learning_rate=args.learning_rate)
      __lowerCamelCase								=						tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
      __lowerCamelCase								=						[tf.keras.metrics.SparseCategoricalAccuracy()]
      model.compile(optimizer=optimizer, loss=loss, metrics=metrics)

      # Training
      if args.do_train:
            # train_results = model.fit(tf_train_dataset, epochs=args.epochs, batch_size=args.train_batch_size)
            __lowerCamelCase								=						time.time()
            __lowerCamelCase								=						fit(
                model,
                loss,
                optimizer,
                tf_train_dataset,
                args.epochs,
                args.per_device_train_batch_size,
                max_steps=args.max_steps,
            )
            __lowerCamelCase								=						time.time() - start_train_time
            logger.info("*** Train ***")
            logger.info(f'''train_runtime = {end_train_time}''')

            __lowerCamelCase								=						os.path.join(args.output_dir, "train_results.txt")

            if not SDP_ENABLED or sdp.rank() == 0:
                  with open(output_eval_file, "w") as writer:
                        logger.info("***** Train results *****")
                        logger.info(train_results)
                        for key, value in train_results.items():
                              logger.info(f'''  {key} = {value}''')
                              writer.write(f'''{key} = {value}\n''')

    # Evaluation
      if args.do_eval and (not SDP_ENABLED or sdp.rank() == 0):
            __lowerCamelCase								=						model.evaluate(tf_test_dataset, batch_size=args.per_device_eval_batch_size, return_dict=True)
            logger.info("*** Evaluate ***")

            __lowerCamelCase								=						os.path.join(args.output_dir, "eval_results.txt")

            with open(output_eval_file, "w") as writer:
                  logger.info("***** Eval results *****")
                  logger.info(result)
                  for key, value in result.items():
                        logger.info(f'''  {key} = {value}''')
                        writer.write(f'''{key} = {value}\n''')

    # Save result
      if SDP_ENABLED:
            if sdp.rank() == 0:
                  model.save_pretrained(args.output_dir)
                  tokenizer.save_pretrained(args.output_dir)
      else:
            model.save_pretrained(args.output_dir)
            tokenizer.save_pretrained(args.output_dir)








"""simple docstring"""

from typing import Dict, Tuple, overload

import torch
import torch.types
from torch import nn

from . import residue_constants as rc
from .rigid_utils import Rigid, Rotation
from .tensor_utils import batched_gather



@overload
def 				a (						__snake_case     :				torch.Tensor,    __snake_case     :				torch.Tensor,    __snake_case     :				None  ):



				'''simple docstring'''





				...



@overload
def 				a (						__snake_case     :				torch.Tensor,    __snake_case     :				torch.Tensor,    __snake_case     :				torch.Tensor  ):



				'''simple docstring'''





				...



def 				a (						__snake_case     :				List[str],    __snake_case     :				Optional[int],    __snake_case     :				str  ):



				'''simple docstring'''





				UpperCAmelCase_					:int      			=      aatype == rc.restype_order['''G''']
				UpperCAmelCase_					:Optional[Any]      			=      rc.atom_order['''CA''']
				UpperCAmelCase_					:Optional[Any]      			=      rc.atom_order['''CB''']
				UpperCAmelCase_					:Optional[Any]      			=      torch.where(
				    is_gly[..., None].expand(*((-1,) * len(is_gly.shape  )),    3  ),    all_atom_positions[..., ca_idx, :],    all_atom_positions[..., cb_idx, :],    )

				if all_atom_masks is not None:
								UpperCAmelCase_					:Any      			=      torch.where(
								    __snake_case,    all_atom_masks[..., ca_idx],    all_atom_masks[..., cb_idx],    )
								return pseudo_beta, pseudo_beta_mask
				else:
								return pseudo_beta



def 				a (						__snake_case     :				torch.Tensor,    __snake_case     :				Dict[str, torch.Tensor]  ):



				'''simple docstring'''





				UpperCAmelCase_					:List[Any]      			=      batched_gather(
				    __snake_case,    batch['''residx_atom37_to_atom14'''],    dim=-2,    no_batch_dims=len(atomaa.shape[:-2]  ),    )

				UpperCAmelCase_					:Optional[Any]      			=      atomaa_data * batch['''atom37_atom_exists'''][..., None]

				return atomaa_data



def 				a (						__snake_case     :				Dict[str, torch.Tensor]  ):



				'''simple docstring'''





				UpperCAmelCase_					:str      			=      template_feats['''template_aatype''']
				UpperCAmelCase_					:Union[str, Any]      			=      template_feats['''template_torsion_angles_sin_cos''']
				UpperCAmelCase_					:Tuple      			=      template_feats['''template_alt_torsion_angles_sin_cos''']
				UpperCAmelCase_					:Tuple      			=      template_feats['''template_torsion_angles_mask''']
				UpperCAmelCase_					:int      			=      torch.cat(
				    [
				        nn.functional.one_hot(__snake_case,    22  ),
				        torsion_angles_sin_cos.reshape(*torsion_angles_sin_cos.shape[:-2],    14  ),
				        alt_torsion_angles_sin_cos.reshape(*alt_torsion_angles_sin_cos.shape[:-2],    14  ),
				        torsion_angles_mask,
				    ],    dim=-1,    )

				return template_angle_feat



def 				a (						__snake_case     :				Dict[str, torch.Tensor],    __snake_case     :				torch.types.Number,    __snake_case     :				torch.types.Number,    __snake_case     :				int,    __snake_case     :				bool = False,    __snake_case     :				float = 1E-20,    __snake_case     :				float = 1E8,    ):



				'''simple docstring'''





				UpperCAmelCase_					:List[str]      			=      batch['''template_pseudo_beta_mask''']
				UpperCAmelCase_					:Any      			=      template_mask[..., None] * template_mask[..., None, :]

				# Compute distogram (this seems to differ slightly from Alg. 5)
				UpperCAmelCase_					:List[str]      			=      batch['''template_pseudo_beta''']
				UpperCAmelCase_					:Any      			=      torch.sum((tpb[..., None, :] - tpb[..., None, :, :]) ** 2,    dim=-1,    keepdim=__snake_case  )
				UpperCAmelCase_					:Any      			=      torch.linspace(__snake_case,    __snake_case,    __snake_case,    device=tpb.device  ) ** 2
				UpperCAmelCase_					:Union[str, Any]      			=      torch.cat([lower[1:], lower.new_tensor([inf]  )],    dim=-1  )
				UpperCAmelCase_					:Tuple      			=      ((dgram > lower) * (dgram < upper)).type(dgram.dtype  )

				UpperCAmelCase_					:List[str]      			=      [dgram, template_mask_ad[..., None]]

				UpperCAmelCase_					:torch.LongTensor      			=      nn.functional.one_hot(
				    batch['''template_aatype'''],    rc.restype_num + 2,    )

				UpperCAmelCase_					:Optional[Any]      			=      batch['''template_aatype'''].shape[-1]
				to_concat.append(aatype_one_hot[..., None, :, :].expand(*aatype_one_hot.shape[:-2],    __snake_case,    -1,    -1  )  )
				to_concat.append(aatype_one_hot[..., None, :].expand(*aatype_one_hot.shape[:-2],    -1,    __snake_case,    -1  )  )

				UpperCAmelCase_			,UpperCAmelCase_			,UpperCAmelCase_					:Optional[int]      			=      [rc.atom_order[a] for a in ['''N''', '''CA''', '''C''']]
				UpperCAmelCase_					:str      			=      Rigid.make_transform_from_reference(
				    n_xyz=batch['''template_all_atom_positions'''][..., n, :],    ca_xyz=batch['''template_all_atom_positions'''][..., ca, :],    c_xyz=batch['''template_all_atom_positions'''][..., c, :],    eps=__snake_case,    )
				UpperCAmelCase_					:Dict      			=      rigids.get_trans()[..., None, :, :]
				UpperCAmelCase_					:Dict      			=      rigids[..., None].invert_apply(__snake_case  )

				UpperCAmelCase_					:Optional[int]      			=      torch.rsqrt(eps + torch.sum(rigid_vec**2,    dim=-1  )  )

				UpperCAmelCase_					:Optional[int]      			=      batch['''template_all_atom_mask''']
				UpperCAmelCase_					:Tuple      			=      t_aa_masks[..., n] * t_aa_masks[..., ca] * t_aa_masks[..., c]
				UpperCAmelCase_					:int      			=      template_mask[..., None] * template_mask[..., None, :]

				UpperCAmelCase_					:str      			=      inv_distance_scalar * template_mask_ad
				UpperCAmelCase_					:List[Any]      			=      rigid_vec * inv_distance_scalar[..., None]

				if not use_unit_vector:
								UpperCAmelCase_					:List[Any]      			=      unit_vector * 0.0

				to_concat.extend(torch.unbind(unit_vector[..., None, :],    dim=-1  )  )
				to_concat.append(template_mask_ad[..., None]  )

				UpperCAmelCase_					:Optional[int]      			=      torch.cat(__snake_case,    dim=-1  )
				UpperCAmelCase_					:Optional[Any]      			=      act * template_mask_ad[..., None]

				return act



def 				a (						__snake_case     :				Dict[str, torch.Tensor]  ):



				'''simple docstring'''





				UpperCAmelCase_					:torch.LongTensor      			=      nn.functional.one_hot(batch['''extra_msa'''],    23  )
				UpperCAmelCase_					:Union[str, Any]      			=      [
				    msa_ahot,
				    batch['''extra_has_deletion'''].unsqueeze(-1  ),
				    batch['''extra_deletion_value'''].unsqueeze(-1  ),
				]
				return torch.cat(__snake_case,    dim=-1  )



def 				a (						__snake_case     :				Rigid,    __snake_case     :				torch.Tensor,    __snake_case     :				torch.Tensor,    __snake_case     :				torch.Tensor,    ):



				'''simple docstring'''





				UpperCAmelCase_					:Optional[int]      			=      rrgdf[aatype, ...]

				# [*, N, 8] transformations, i.e.
				#   One [*, N, 8, 3, 3] rotation matrix and
				#   One [*, N, 8, 3]    translation matrix
				UpperCAmelCase_					:Optional[int]      			=      r.from_tensor_axa(__snake_case  )

				UpperCAmelCase_					:List[Any]      			=      alpha.new_zeros((*((1,) * len(alpha.shape[:-1]  )), 2)  )
				UpperCAmelCase_					:List[Any]      			=      1

				# [*, N, 8, 2]
				UpperCAmelCase_					:Tuple      			=      torch.cat([bb_rot.expand(*alpha.shape[:-2],    -1,    -1  ), alpha],    dim=-2  )

				# [*, N, 8, 3, 3]
				# Produces rotation matrices of the form:
				# [
				#   [1, 0  , 0  ],
				#   [0, a_2,-a_1],
				#   [0, a_1, a_2]
				# ]
				# This follows the original code rather than the supplement, which uses
				# different indices.

				UpperCAmelCase_					:Any      			=      alpha.new_zeros(default_r.get_rots().get_rot_mats().shape  )
				UpperCAmelCase_					:Optional[Any]      			=      1
				UpperCAmelCase_					:int      			=      alpha[..., 1]
				UpperCAmelCase_					:Optional[int]      			=      -alpha[..., 0]
				UpperCAmelCase_					:int      			=      alpha

				UpperCAmelCase_					:List[str]      			=      default_r.compose(Rigid(Rotation(rot_mats=__snake_case  ),    __snake_case  )  )

				UpperCAmelCase_					:List[str]      			=      all_frames[..., 5]
				UpperCAmelCase_					:str      			=      all_frames[..., 6]
				UpperCAmelCase_					:List[Any]      			=      all_frames[..., 7]

				UpperCAmelCase_					:Tuple      			=      all_frames[..., 4]
				UpperCAmelCase_					:Optional[Any]      			=      chia_frame_to_bb.compose(__snake_case  )
				UpperCAmelCase_					:Optional[Any]      			=      chia_frame_to_bb.compose(__snake_case  )
				UpperCAmelCase_					:Union[str, Any]      			=      chia_frame_to_bb.compose(__snake_case  )

				UpperCAmelCase_					:List[str]      			=      Rigid.cat(
				    [
				        all_frames[..., :5],
				        chia_frame_to_bb.unsqueeze(-1  ),
				        chia_frame_to_bb.unsqueeze(-1  ),
				        chia_frame_to_bb.unsqueeze(-1  ),
				    ],    dim=-1,    )

				UpperCAmelCase_					:Any      			=      r[..., None].compose(__snake_case  )

				return all_frames_to_global





def 				a (						__snake_case     :				Rigid,    __snake_case     :				torch.Tensor,    __snake_case     :				torch.Tensor,    __snake_case     :				torch.Tensor,    __snake_case     :				torch.Tensor,    __snake_case     :				torch.Tensor,    ):



				'''simple docstring'''





				UpperCAmelCase_					:Optional[Any]      			=      group_idx[aatype, ...]

				# [*, N, 14, 8]
				UpperCAmelCase_					:torch.LongTensor      			=      nn.functional.one_hot(
				    __snake_case,    num_classes=default_frames.shape[-3],    )

				# [*, N, 14, 8]
				UpperCAmelCase_					:Dict      			=      r[..., None, :] * group_mask_one_hot

				# [*, N, 14]
				UpperCAmelCase_					:List[Any]      			=      t_atoms_to_global.map_tensor_fn(lambda __snake_case  : torch.sum(__snake_case,    dim=-1  )  )

				# [*, N, 14, 1]
				UpperCAmelCase_					:List[Any]      			=      atom_mask[aatype, ...].unsqueeze(-1  )

				# [*, N, 14, 3]
				UpperCAmelCase_					:List[Any]      			=      lit_positions[aatype, ...]
				UpperCAmelCase_					:Optional[int]      			=      t_atoms_to_global.apply(__snake_case  )
				UpperCAmelCase_					:int      			=      pred_positions * atom_mask

				return pred_positions


