from copy import deepcopy

import torch
import torch.nn.functional as F
from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader

from accelerate.accelerator import Accelerator
from accelerate.state import GradientState
from accelerate.test_utils import RegressionDataset, RegressionModel
from accelerate.utils import DistributedType, is_torch_version, set_seed







def  __UpperCamelCase    (						A ,							A ,							A ,							A							):
     for param, grad_param in zip(model_a.parameters() ,							model_b.parameters()							):
          if not param.requires_grad:
               continue
          if not did_step:
               # Grads should not be in sync
               assert (
                   torch.allclose(param.grad ,							grad_param.grad							) is False
               ), f"Gradients in sync when they should not be at iteration {iteration}:\nmodel_a grad ({param.grad}) == model_b grad ({grad_param.grad})"
          else:
               # Grads should be in sync
               assert (
                   torch.allclose(param.grad ,							grad_param.grad							) is True
               ), f"Gradients not in sync when they should be at iteration {iteration}:\nmodel_a grad ({param.grad}) != model_b grad ({grad_param.grad})"







def  __UpperCamelCase    (						A ,							A ,							A ,							A ,							A=True							):
     model.train()
     UpperCamelCase__										=      model(lowerCAmelCase__							)
     UpperCamelCase__										=      F.mse_loss(lowerCAmelCase__ ,							target.to(output.device							)							)
     if not do_backward:
          loss /= accelerator.gradient_accumulation_steps
          loss.backward()
     else:
          accelerator.backward(lowerCAmelCase__							)







def  __UpperCamelCase    (						A ,							A=False							):
     set_seed(42							)
     UpperCamelCase__										=      RegressionModel()
     UpperCamelCase__										=      deepcopy(lowerCAmelCase__							)
     UpperCamelCase__										=      RegressionDataset(length=80							)
     UpperCamelCase__										=      DataLoader(lowerCAmelCase__ ,							batch_size=16							)
     model.to(accelerator.device							)
     if sched:
          UpperCamelCase__										=      AdamW(params=model.parameters() ,							lr=1e-3							)
          UpperCamelCase__										=      AdamW(params=ddp_model.parameters() ,							lr=1e-3							)
          UpperCamelCase__										=      LambdaLR(lowerCAmelCase__ ,							lr_lambda=lambda A							: epoch**0.65							)
          UpperCamelCase__										=      LambdaLR(lowerCAmelCase__ ,							lr_lambda=lambda A							: epoch**0.65							)
     # Make a copy of `model`
     if sched:
          UpperCamelCase__						,       UpperCamelCase__						,       UpperCamelCase__						,       UpperCamelCase__										=      accelerator.prepare(lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__							)
     else:
          UpperCamelCase__						,       UpperCamelCase__										=      accelerator.prepare(lowerCAmelCase__ ,							lowerCAmelCase__							)
     if sched:
          return (model, opt, sched, dataloader, ddp_model, ddp_opt, ddp_sched)
     return model, ddp_model, dataloader







def  __UpperCamelCase    (						A							):
     # Test when on a single CPU or GPU that the context manager does nothing
     UpperCamelCase__						,       UpperCamelCase__						,       UpperCamelCase__										=      get_training_setup(lowerCAmelCase__							)
     # Use a single batch
     UpperCamelCase__						,       UpperCamelCase__										=      next(iter(lowerCAmelCase__							)							).values()
     for iteration in range(3							):
          # Gather the distributed inputs and targs for the base model
          UpperCamelCase__						,       UpperCamelCase__										=      accelerator.gather((ddp_input, ddp_target)							)
          UpperCamelCase__						,       UpperCamelCase__										=      input.to(accelerator.device							), target.to(accelerator.device							)
          # Perform our initial ground truth step in non "DDP"
          step_model(lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__							)
          # Do "gradient accumulation" (noop)
          if iteration % 2 == 0:
               # Accumulate grads locally
               with accelerator.no_sync(lowerCAmelCase__							):
                    step_model(lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__							)
          else:
               # Sync grads
               step_model(lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__							)

          # Since `no_sync` is a noop, `ddp_model` and `model` grads should always be in sync
          check_model_parameters(lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__							)
          for param, ddp_param in zip(model.parameters() ,							ddp_model.parameters()							):
               if not param.requires_grad:
                    continue
               assert torch.allclose(
                   param.grad ,							ddp_param.grad							), f"Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})"

          # Shuffle ddp_input on each iteration
          torch.manual_seed(1337 + iteration							)
          UpperCamelCase__										=      ddp_input[torch.randperm(len(lowerCAmelCase__							)							)]







def  __UpperCamelCase    (						A							):
     # Test on distributed setup that context manager behaves properly
     UpperCamelCase__						,       UpperCamelCase__						,       UpperCamelCase__										=      get_training_setup(lowerCAmelCase__							)
     # Use a single batch
     UpperCamelCase__						,       UpperCamelCase__										=      next(iter(lowerCAmelCase__							)							).values()
     for iteration in range(3							):
          # Gather the distributed inputs and targs for the base model
          UpperCamelCase__						,       UpperCamelCase__										=      accelerator.gather((ddp_input, ddp_target)							)
          UpperCamelCase__						,       UpperCamelCase__										=      input.to(accelerator.device							), target.to(accelerator.device							)
          # Perform our initial ground truth step in non "DDP"
          step_model(lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__							)
          # Do "gradient accumulation" (noop)
          if iteration % 2 == 0:
               # Accumulate grads locally
               with accelerator.no_sync(lowerCAmelCase__							):
                    step_model(lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__							)
          else:
               # Sync grads
               step_model(lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__							)

          # DDP model and model should only be in sync when not (iteration % 2 == 0)
          for param, ddp_param in zip(model.parameters() ,							ddp_model.parameters()							):
               if not param.requires_grad:
                    continue
               if iteration % 2 == 0:
                    # Grads should not be in sync
                    assert (
                        torch.allclose(param.grad ,							ddp_param.grad							) is False
                    ), f"Gradients in sync when they should not be:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})"
               else:
                    # Grads should be in sync
                    assert (
                        torch.allclose(param.grad ,							ddp_param.grad							) is True
                    ), f"Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})"

        # Shuffle ddp_input on each iteration
          torch.manual_seed(1337 + iteration							)
          UpperCamelCase__										=      ddp_input[torch.randperm(len(lowerCAmelCase__							)							)]







def  __UpperCamelCase    (						A=False ,							A=False							):
     UpperCamelCase__										=      Accelerator(
         split_batches=lowerCAmelCase__ ,							dispatch_batches=lowerCAmelCase__ ,							gradient_accumulation_steps=2							)
     # Test that context manager behaves properly
     UpperCamelCase__						,       UpperCamelCase__						,       UpperCamelCase__										=      get_training_setup(lowerCAmelCase__							)
     for iteration, batch in enumerate(lowerCAmelCase__							):
          UpperCamelCase__						,       UpperCamelCase__										=      batch.values()
          # Gather the distributed inputs and targs for the base model
          UpperCamelCase__						,       UpperCamelCase__										=      accelerator.gather((ddp_input, ddp_target)							)
          UpperCamelCase__						,       UpperCamelCase__										=      input.to(accelerator.device							), target.to(accelerator.device							)
          # Perform our initial ground truth step in non "DDP"
          step_model(lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__							)
          # Do "gradient accumulation" (noop)
          with accelerator.accumulate(lowerCAmelCase__							):
               step_model(lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__							)

          # DDP model and model should only be in sync when not (iteration % 2 == 0)
          for param, ddp_param in zip(model.parameters() ,							ddp_model.parameters()							):
               if not param.requires_grad:
                    continue
               if ((iteration + 1) % 2 == 0) or (iteration == len(lowerCAmelCase__							) - 1):
                    # Grads should be in sync
                    assert (
                        torch.allclose(param.grad ,							ddp_param.grad							) is True
                    ), f"Gradients not in sync when they should be at iteration {iteration}:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})"
               else:
                    # Grads should not be in sync
                    assert (
                        torch.allclose(param.grad ,							ddp_param.grad							) is False
                    ), f"Gradients in sync when they should not be at iteration {iteration}:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})"

        # Shuffle ddp_input on each iteration
          torch.manual_seed(1337 + iteration							)
          UpperCamelCase__										=      ddp_input[torch.randperm(len(lowerCAmelCase__							)							)]
     GradientState._reset_state()







def  __UpperCamelCase    (						A=False ,							A=False							):
     UpperCamelCase__										=      Accelerator(
         split_batches=lowerCAmelCase__ ,							dispatch_batches=lowerCAmelCase__ ,							gradient_accumulation_steps=2							)
     # Test that context manager behaves properly
     UpperCamelCase__						,       UpperCamelCase__						,       UpperCamelCase__						,       UpperCamelCase__						,       UpperCamelCase__						,       UpperCamelCase__						,       UpperCamelCase__										=      get_training_setup(lowerCAmelCase__ ,							lowerCAmelCase__							)
     for iteration, batch in enumerate(lowerCAmelCase__							):
          UpperCamelCase__						,       UpperCamelCase__										=      batch.values()
          # Gather the distributed inputs and targs for the base model
          UpperCamelCase__						,       UpperCamelCase__										=      accelerator.gather((ddp_input, ddp_target)							)
          UpperCamelCase__						,       UpperCamelCase__										=      input.to(accelerator.device							), target.to(accelerator.device							)
          # Perform our initial ground truth step in non "DDP"
          model.train()
          ddp_model.train()
          step_model(lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__							)
          opt.step()

          if ((iteration + 1) % 2 == 0) or ((iteration + 1) == len(lowerCAmelCase__							)):
               if split_batches:
                    sched.step()
               else:
                    for _ in range(accelerator.num_processes							):
                         sched.step()
          opt.zero_grad()
          # Perform gradient accumulation under wrapper
          with accelerator.accumulate(lowerCAmelCase__							):
               step_model(lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__							)
               ddp_opt.step()
               ddp_sched.step()
               ddp_opt.zero_grad()

          # Learning rates should be the same
          assert (
              opt.param_groups[0]["lr"] == ddp_opt.param_groups[0]["lr"]
          ), f"Learning rates found in each optimizer did not align\nopt: {opt.param_groups[0]['lr']}\nDDP opt: {ddp_opt.param_groups[0]['lr']}\n"
          UpperCamelCase__										=      (((iteration + 1) % 2) == 0) or ((iteration + 1) == len(lowerCAmelCase__							))
          if accelerator.num_processes > 1:
               check_model_parameters(lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__ ,							lowerCAmelCase__							)
          # Shuffle ddp_input on each iteration
          torch.manual_seed(1337 + iteration							)
     GradientState._reset_state()







def  __UpperCamelCase    (						):
     UpperCamelCase__										=      Accelerator()

     UpperCamelCase__										=      RegressionDataset(length=80							)
     UpperCamelCase__										=      DataLoader(lowerCAmelCase__ ,							batch_size=16							)
     UpperCamelCase__										=      RegressionDataset(length=96							)
     UpperCamelCase__										=      DataLoader(lowerCAmelCase__ ,							batch_size=16							)
     UpperCamelCase__						,       UpperCamelCase__										=      accelerator.prepare(lowerCAmelCase__ ,							lowerCAmelCase__							)
     assert accelerator.gradient_state.active_dataloader is None
     for iteration, _ in enumerate(lowerCAmelCase__							):
          assert id(accelerator.gradient_state.active_dataloader							) == id(lowerCAmelCase__							)
          if iteration < len(lowerCAmelCase__							) - 1:
               assert not accelerator.gradient_state.end_of_dataloader
               if iteration == 1:
                    for batch_num, _ in enumerate(lowerCAmelCase__							):
                         assert id(accelerator.gradient_state.active_dataloader							) == id(lowerCAmelCase__							)
                         if batch_num < len(lowerCAmelCase__							) - 1:
                              assert not accelerator.gradient_state.end_of_dataloader
                         else:
                              assert accelerator.gradient_state.end_of_dataloader
          else:
               assert accelerator.gradient_state.end_of_dataloader
     assert accelerator.gradient_state.active_dataloader is None







def  __UpperCamelCase    (						):
     UpperCamelCase__										=      Accelerator()
     UpperCamelCase__										=      accelerator.state
     if state.local_process_index == 0:
          print('''**Test `accumulate` gradient accumulation with dataloader break**'''							)
     test_dataloader_break()
     if state.distributed_type == DistributedType.NO:
          if state.local_process_index == 0:
               print('''**Test NOOP `no_sync` context manager**'''							)
          test_noop_sync(lowerCAmelCase__							)
     if state.distributed_type in (DistributedType.MULTI_GPU, DistributedType.MULTI_CPU):
          if state.local_process_index == 0:
               print('''**Test Distributed `no_sync` context manager**'''							)
          test_distributed_sync(lowerCAmelCase__							)
     if state.distributed_type == DistributedType.MULTI_GPU:
          for split_batch in [True, False]:
               for dispatch_batches in [True, False]:
                    if state.local_process_index == 0:
                         print(
                             '''**Test `accumulate` gradient accumulation, ''' ,							f"`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**" ,							)
                    test_gradient_accumulation(lowerCAmelCase__ ,							lowerCAmelCase__							)

    # Currently will break on torch 2.0 +, need to investigate why
     if is_torch_version('''<''' ,							'''2.0'''							) or state.distributed_type == DistributedType.NO:
          if state.local_process_index == 0:
               print(
                   '''**Test `accumulate` gradient accumulation with optimizer and scheduler, ''' ,							'''`split_batches=False`, `dispatch_batches=False`**''' ,							)
          test_gradient_accumulation_with_opt_and_scheduler()
          if state.distributed_type == DistributedType.MULTI_GPU:
               for split_batch in [True, False]:
                    for dispatch_batches in [True, False]:
                         if not split_batch and not dispatch_batches:
                              continue
                         if state.local_process_index == 0:
                              print(
                                  '''**Test `accumulate` gradient accumulation with optimizer and scheduler, ''' ,							f"`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**" ,							)
                         test_gradient_accumulation_with_opt_and_scheduler(lowerCAmelCase__ ,							lowerCAmelCase__							)







def  __UpperCamelCase    (						A							):
     # For xla_spawn (TPUs)
     main()


if __name__ == "__main__":
    main()




import argparse
import json
import os

import tensorstore as ts
import torch
from flax import serialization
from flax.traverse_util import flatten_dict, unflatten_dict
from tensorflow.io import gfile

from transformers.modeling_utils import dtype_byte_size
from transformers.models.switch_transformers.convert_switch_transformers_original_flax_checkpoint_to_pytorch import (
    rename_keys,
)
from transformers.utils import WEIGHTS_INDEX_NAME, WEIGHTS_NAME
from transformers.utils.hub import convert_file_size_to_int







def  __UpperCamelCase    (						A ,							A							):
     if flax_key_tuple[-1] == "kernel" and flax_tensor.ndim == 3:
          # expert layer
          UpperCamelCase__										=      flax_key_tuple[:-1] + ('''weight''',)
          UpperCamelCase__										=      torch.permute(A ,							(0, 2, 1)							)
     elif flax_key_tuple[-1] == "kernel" and ".".join(A							):
          # linear layer
          UpperCamelCase__										=      flax_key_tuple[:-1] + ('''weight''',)
          UpperCamelCase__										=      flax_tensor.T
     elif flax_key_tuple[-1] in ["scale", "embedding"]:
          UpperCamelCase__										=      flax_key_tuple[:-1] + ('''weight''',)

     return flax_key_tuple, flax_tensor







def  __UpperCamelCase    (						A ,							A ,							A							):
     if "metadata" in layer:
          UpperCamelCase__										=      layer.split('''metadata'''							)
          UpperCamelCase__										=      ''''''.join(split_layer[0]							)[:-1]
          UpperCamelCase__										=      [tuple(('''metadata''' + split_layer[1]).split('''/'''							)							)]
     elif "kvstore" in layer:
          UpperCamelCase__										=      layer.split('''kvstore'''							)
          UpperCamelCase__										=      ''''''.join(split_layer[0]							)[:-1]
          UpperCamelCase__										=      [tuple(('''kvstore''' + split_layer[1]).split('''/'''							)							)]

     else:
          UpperCamelCase__										=      layer.split('''/'''							)
          UpperCamelCase__										=      '''/'''.join(split_layer[:-1]							)
          UpperCamelCase__										=      (split_layer[-1],)

     if "kvstore/path" in layer:
          UpperCamelCase__										=      f"{switch_checkpoint_path}/{checkpoint_info[layer]}"
     elif "kvstore/driver" in layer:
          UpperCamelCase__										=      '''file'''
     else:
          UpperCamelCase__										=      checkpoint_info[layer]

     return curr_real_layer_name, split_layer, content







def  __UpperCamelCase    (						A ,							A							):
     UpperCamelCase__										=      rename_keys(A							)
     UpperCamelCase__										=      {}
     for k, v in current_block.items():
          UpperCamelCase__										=      v
     UpperCamelCase__										=      new_current_block
     torch.save(A ,							A							)







def  __UpperCamelCase    (						A ,							A ,							A ,							A ,							A = WEIGHTS_NAME							):
     UpperCamelCase__										=      convert_file_size_to_int(A							)
     UpperCamelCase__										=      []
     UpperCamelCase__										=      {}
     UpperCamelCase__										=      0
     UpperCamelCase__										=      0

     os.makedirs(A ,							exist_ok=A							)
     with gfile.GFile(switch_checkpoint_path + '''/checkpoint''' ,							'''rb'''							) as fp:
          UpperCamelCase__										=      serialization.msgpack_restore(fp.read()							)['''optimizer''']['''target''']
          UpperCamelCase__										=      flatten_dict(A ,							sep='''/'''							)

     UpperCamelCase__										=      {}
     for layer in checkpoint_info.keys():
          UpperCamelCase__						,       UpperCamelCase__						,       UpperCamelCase__										=      get_key_and_tensorstore_dict(
              A ,							A ,							A							)
          if curr_real_layer_name in all_layers:
               UpperCamelCase__										=      content
          else:
               UpperCamelCase__										=      {split_layer[-1]: content}

     for key in all_layers.keys():
          # open tensorstore file
          UpperCamelCase__										=      ts.open(unflatten_dict(all_layers[key]							)							).result().read().result()
          UpperCamelCase__										=      torch.tensor(A							)
          UpperCamelCase__										=      raw_weights.numel() * dtype_byte_size(raw_weights.dtype							)

          # use the renaming pattern from the small conversion scripts
          UpperCamelCase__						,       UpperCamelCase__										=      rename_base_flax_keys(tuple(key.split('''/'''							)							) ,							A							)
          UpperCamelCase__										=      '''/'''.join(A							)

          # If this weight is going to tip up over the maximal size, we split.
          if current_block_size + weight_size > max_shard_size:
               UpperCamelCase__										=      os.path.join(
                   A ,							weights_name.replace('''.bin''' ,							f"-{len(A							)+1:05d}-of-???.bin"							)							)
               rename_and_save_block(A ,							A							)
               sharded_state_dicts.append(current_block.keys()							)
               del current_block
               UpperCamelCase__										=      {}
               UpperCamelCase__										=      0

          UpperCamelCase__										=      raw_weights.to(getattr(A ,							A							)							)
          current_block_size += weight_size
          total_size += weight_size

     # Add the last block
     UpperCamelCase__										=      os.path.join(A ,							weights_name.replace('''.bin''' ,							f"-{len(A							)+1:05d}-of-???.bin"							)							)
     rename_and_save_block(A ,							A							)
     sharded_state_dicts.append(current_block.keys()							)

     # If we only have one shard, we return it
     if len(A							) == 1:
          return {weights_name: sharded_state_dicts[0]}, None

     # Otherwise, let's build the index
     UpperCamelCase__										=      {}
     UpperCamelCase__										=      {}
     for idx, shard in enumerate(A							):
          UpperCamelCase__										=      weights_name.replace(
              '''.bin''' ,							f"-{idx+1:05d}-of-{len(A							):05d}.bin"							)  # len(sharded_state_dicts):05d}
          UpperCamelCase__										=      os.path.join(A ,							weights_name.replace('''.bin''' ,							f"-{idx+1:05d}-of-???.bin"							)							)
          os.rename(A ,							os.path.join(A ,							A							)							)
          UpperCamelCase__										=      shard
          for key in shard:
               UpperCamelCase__										=      shard_file

    # Add the metadata
     UpperCamelCase__										=      {'''total_size''': total_size}
     UpperCamelCase__										=      {'''metadata''': metadata, '''weight_map''': weight_map}

     with open(os.path.join(A ,							A							) ,							'''w''' ,							encoding='''utf-8'''							) as f:
          UpperCamelCase__										=      json.dumps(A ,							indent=2 ,							sort_keys=A							) + '''\n'''
          f.write(A							)

     return metadata, index


if __name__ == "__main__":
     __magic_name__								=argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         '''--switch_t5x_checkpoint_path''',
         default='''/mnt/disks/disk_switch/original_checkpoints/switch-xxl-128/checkpoint_634600''',
         type=str,
         required=False,
         help='''Path to a directory containing a folder per layer. Follows the original Google format.''',
     )
     parser.add_argument('''--max_shard_size''', default='''10GB''', required=False, help='''Max shard size''')
     parser.add_argument('''--dtype''', default='''bfloat16''', type=str, required=False, help='''dtype of the saved model''')
     parser.add_argument(
         '''--pytorch_dump_folder_path''',
         default='''/mnt/disks/disk_switch/original_checkpoints/switch-xxl-128-converted''',
         type=str,
         required=False,
         help='''Path to the output pytorch model.''',
     )
     __magic_name__								=parser.parse_args()
     shard_on_the_fly(
         args.switch_tax_checkpoint_path,
         args.pytorch_dump_folder_path,
         args.max_shard_size,
         args.dtype,
     )







def  __UpperCamelCase    (						):
     from transformers import SwitchTransformersConfig, SwitchTransformersForConditionalGeneration, TaTokenizer

     UpperCamelCase__										=      SwitchTransformersConfig.from_pretrained('''google/switch-base-8'''							)
     config.save_pretrained('''/home/arthur_huggingface_co/transformers/switch_converted'''							)
     UpperCamelCase__										=      SwitchTransformersForConditionalGeneration.from_pretrained(
         '''/home/arthur_huggingface_co/transformers/switch_converted''' ,							device_map='''auto'''							)

     UpperCamelCase__										=      TaTokenizer.from_pretrained('''t5-small'''							)
     UpperCamelCase__										=      '''A <extra_id_0> walks into a bar a orders a <extra_id_1> with <extra_id_2> pinch of <extra_id_3>.'''

     UpperCamelCase__										=      tokenizer(A ,							return_tensors='''pt'''							).input_ids
     UpperCamelCase__										=      model.generate(A ,							decoder_start_token_id=0							)
     print(tokenizer.decode(out[0]							)							)



