


import argparse
import json
import os

import evaluate
import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType
from accelerate.utils.deepspeed import DummyOptim, DummyScheduler


_SCREAMING_SNAKE_CASE      :			Optional[int]      				=   16
_SCREAMING_SNAKE_CASE      :			int      				=   32
def       UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_ = 16							,UpperCamelCase_ = "bert-base-cased"    ):



     """simple docstring"""






     snake_case					=    AutoTokenizer.from_pretrained(UpperCamelCase_    )
     snake_case					=    load_dataset('''glue'''							,'''mrpc'''    )

     def tokenize_function(UpperCamelCase_    ):
          # max_length=None => use the model max length (it's actually the default)
          snake_case					=    tokenizer(examples['''sentence1''']							,examples['''sentence2''']							,truncation=UpperCamelCase_							,max_length=UpperCamelCase_    )
          return outputs

     # Apply the method we just defined to all the examples in all the splits of the dataset
     snake_case					=    datasets.map(
         UpperCamelCase_							,batched=UpperCamelCase_							,remove_columns=['''idx''', '''sentence1''', '''sentence2''']							,load_from_cache_file=UpperCamelCase_    )

     # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
     # transformers library
     snake_case					=    tokenized_datasets.rename_column('''label'''							,'''labels'''    )

     def collate_fn(UpperCamelCase_    ):
          # On TPU it's best to pad everything to the same length or training will be very slow.
          if accelerator.distributed_type == DistributedType.TPU:
               return tokenizer.pad(UpperCamelCase_							,padding='''max_length'''							,max_length=1_28							,return_tensors='''pt'''    )
          return tokenizer.pad(UpperCamelCase_							,padding='''longest'''							,return_tensors='''pt'''    )

     # Instantiate dataloaders.
     snake_case					=    DataLoader(
         tokenized_datasets['''train''']							,shuffle=UpperCamelCase_							,collate_fn=UpperCamelCase_							,batch_size=UpperCamelCase_    )
     snake_case					=    DataLoader(
         tokenized_datasets['''validation''']							,shuffle=UpperCamelCase_							,collate_fn=UpperCamelCase_							,batch_size=UpperCamelCase_    )

     return train_dataloader, eval_dataloader
def       UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_    ):



     """simple docstring"""






     model.eval()
     snake_case					=    0
     for step, batch in enumerate(UpperCamelCase_    ):
          # We could avoid this line since we set the accelerator with `device_placement=True`.
          batch.to(accelerator.device    )
          with torch.no_grad():
               snake_case					=    model(**UpperCamelCase_    )
          snake_case					=    outputs.logits.argmax(dim=-1    )
          # It is slightly faster to call this once, than multiple times
          snake_case			,  snake_case					=    accelerator.gather(
              (predictions, batch['''labels'''])    )  # If we are in a multiprocess environment, the last batch has duplicates
          if accelerator.use_distributed:
               if step == len(UpperCamelCase_    ) - 1:
                    snake_case					=    predictions[: len(eval_dataloader.dataset    ) - samples_seen]
                    snake_case					=    references[: len(eval_dataloader.dataset    ) - samples_seen]
               else:
                    samples_seen += references.shape[0]
          metric.add_batch(
              predictions=UpperCamelCase_							,references=UpperCamelCase_							,)

     snake_case					=    metric.compute()
     return eval_metric["accuracy"]
def       UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_    ):



     """simple docstring"""






     snake_case					=    Accelerator()

     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     snake_case					=    config['''lr''']
     snake_case					=    int(config['''num_epochs''']    )
     snake_case					=    int(config['''seed''']    )
     snake_case					=    int(config['''batch_size''']    )
     snake_case					=    args.model_name_or_path

     set_seed(UpperCamelCase_    )
     snake_case			,  snake_case					=    get_dataloaders(UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_    )

     # Instantiate the model (we build the model here so that the seed also control new weights initialization)
     snake_case					=    AutoModelForSequenceClassification.from_pretrained(UpperCamelCase_							,return_dict=UpperCamelCase_    )

     # Instantiate optimizer
     snake_case					=    (
         AdamW
         if accelerator.state.deepspeed_plugin is None
         or '''optimizer''' not in accelerator.state.deepspeed_plugin.deepspeed_config
         else DummyOptim
     )
     snake_case					=    optimizer_cls(params=model.parameters()							,lr=UpperCamelCase_    )

     if accelerator.state.deepspeed_plugin is not None:
          snake_case					=    accelerator.state.deepspeed_plugin.deepspeed_config[
              '''gradient_accumulation_steps'''
          ]
     else:
          snake_case					=    1
     snake_case					=    (len(UpperCamelCase_    ) * num_epochs) // gradient_accumulation_steps

     # Instantiate scheduler
     if (
         accelerator.state.deepspeed_plugin is None
         or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
     ):
          snake_case					=    get_linear_schedule_with_warmup(
              optimizer=UpperCamelCase_							,num_warmup_steps=0							,num_training_steps=UpperCamelCase_							,)
     else:
          snake_case					=    DummyScheduler(UpperCamelCase_							,total_num_steps=UpperCamelCase_							,warmup_num_steps=0    )

     # Prepare everything
     # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
     # prepare method.
     snake_case			,  snake_case			,  snake_case			,  snake_case			,  snake_case					=    accelerator.prepare(
         UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_    )

     # We need to keep track of how many total steps we have iterated over
     snake_case					=    0
     # We also need to keep track of the stating epoch so files are named properly
     snake_case					=    0
     snake_case					=    evaluate.load('''glue'''							,'''mrpc'''    )
     snake_case					=    num_epochs

     if args.partial_train_epoch is not None:
          snake_case					=    args.partial_train_epoch

     if args.resume_from_checkpoint:
          accelerator.load_state(args.resume_from_checkpoint    )
          snake_case					=    args.resume_from_checkpoint.split('''epoch_'''    )[1]
          snake_case					=    ''''''
          for char in epoch_string:
               if char.isdigit():
                    state_epoch_num += char
               else:
                    break
          snake_case					=    int(UpperCamelCase_    ) + 1
          snake_case					=    evaluation_loop(UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_    )
          accelerator.print('''resumed checkpoint performance:'''							,UpperCamelCase_    )
          accelerator.print('''resumed checkpoint\'s scheduler\'s lr:'''							,lr_scheduler.get_lr()[0]    )
          accelerator.print('''resumed optimizers\'s lr:'''							,optimizer.param_groups[0]['''lr''']    )
          with open(os.path.join(args.output_dir							,F'''state_{starting_epoch-1}.json'''    )							,'''r'''    ) as f:
               snake_case					=    json.load(UpperCamelCase_    )
               assert resumed_state["accuracy"] == accuracy, "Accuracy mismatch, loading from checkpoint failed"
               assert (
                   resumed_state["lr"] == lr_scheduler.get_lr()[0]
               ), "Scheduler learning rate mismatch, loading from checkpoint failed"
               assert (
                   resumed_state["optimizer_lr"] == optimizer.param_groups[0]["lr"]
               ), "Optimizer learning rate mismatch, loading from checkpoint failed"
               assert resumed_state["epoch"] == starting_epoch - 1, "Epoch mismatch, loading from checkpoint failed"
               return

    # Now we train the model
     snake_case					=    {}
     for epoch in range(UpperCamelCase_							,UpperCamelCase_    ):
          model.train()
          for step, batch in enumerate(UpperCamelCase_    ):
               snake_case					=    model(**UpperCamelCase_    )
               snake_case					=    outputs.loss
               snake_case					=    loss / gradient_accumulation_steps
               accelerator.backward(UpperCamelCase_    )
               if step % gradient_accumulation_steps == 0:
                    optimizer.step()
                    lr_scheduler.step()
                    optimizer.zero_grad()

               overall_step += 1
          snake_case					=    F'''epoch_{epoch}'''
          snake_case					=    os.path.join(args.output_dir							,UpperCamelCase_    )
          accelerator.save_state(UpperCamelCase_    )
          snake_case					=    evaluation_loop(UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_    )
          snake_case					=    accuracy
          snake_case					=    lr_scheduler.get_lr()[0]
          snake_case					=    optimizer.param_groups[0]['''lr''']
          snake_case					=    epoch
          snake_case					=    overall_step
          accelerator.print(F'''epoch {epoch}:'''							,UpperCamelCase_    )

          accelerator.wait_for_everyone()
          if accelerator.is_main_process:
               with open(os.path.join(args.output_dir							,F'''state_{epoch}.json'''    )							,'''w'''    ) as f:
                    json.dump(UpperCamelCase_							,UpperCamelCase_    )
def       UpperCAmelCase__						():



     """simple docstring"""






     snake_case					=    argparse.ArgumentParser(description='''Simple example of training script tracking peak GPU memory usage.'''    )
     parser.add_argument(
         '''--model_name_or_path'''							,type=UpperCamelCase_							,default='''bert-base-cased'''							,help='''Path to pretrained model or model identifier from huggingface.co/models.'''							,required=UpperCamelCase_							,)
     parser.add_argument(
         '''--output_dir'''							,type=UpperCamelCase_							,default='''.'''							,help='''Optional save directory where all checkpoint folders will be stored. Default is the current working directory.'''							,)
     parser.add_argument(
         '''--resume_from_checkpoint'''							,type=UpperCamelCase_							,default=UpperCamelCase_							,help='''If the training should continue from a checkpoint folder.'''							,)
     parser.add_argument(
         '''--partial_train_epoch'''							,type=UpperCamelCase_							,default=UpperCamelCase_							,help='''If passed, the training will stop after this number of epochs.'''							,)
     parser.add_argument(
         '''--num_epochs'''							,type=UpperCamelCase_							,default=2							,help='''Number of train epochs.'''							,)
     snake_case					=    parser.parse_args()
     snake_case					=    {'''lr''': 2e-5, '''num_epochs''': args.num_epochs, '''seed''': 42, '''batch_size''': 16}

     training_function(UpperCamelCase_							,UpperCamelCase_    )


if __name__ == "__main__":
 main()








import argparse
import copy
def       UpperCAmelCase__						(UpperCamelCase_    ):



     """simple docstring"""






     snake_case					=    {}

     with open(UpperCamelCase_    ) as f:
          for line in f:
               if line.split()[0] not in dict_of_neighbours:
                    snake_case					=    []
                    _list.append([line.split()[1], line.split()[2]]    )
                    snake_case					=    _list
               else:
                    dict_of_neighbours[line.split()[0]].append(
                        [line.split()[1], line.split()[2]]    )
               if line.split()[1] not in dict_of_neighbours:
                    snake_case					=    []
                    _list.append([line.split()[0], line.split()[2]]    )
                    snake_case					=    _list
               else:
                    dict_of_neighbours[line.split()[1]].append(
                        [line.split()[0], line.split()[2]]    )

     return dict_of_neighbours
def       UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_    ):



     """simple docstring"""






     with open(UpperCamelCase_    ) as f:
          snake_case					=    f.read(1    )
     snake_case					=    start_node

     snake_case					=    []

     snake_case					=    start_node

     snake_case					=    0
     while visiting not in first_solution:
          snake_case					=    1_00_00
          for k in dict_of_neighbours[visiting]:
               if int(k[1]    ) < int(UpperCamelCase_    ) and k[0] not in first_solution:
                    snake_case					=    k[1]
                    snake_case					=    k[0]

          first_solution.append(UpperCamelCase_    )
          snake_case					=    distance_of_first_solution + int(UpperCamelCase_    )
          snake_case					=    best_node

     first_solution.append(UpperCamelCase_    )

     snake_case					=    0
     for k in dict_of_neighbours[first_solution[-2]]:
          if k[0] == start_node:
               break
          position += 1

     snake_case					=    (
         distance_of_first_solution
         + int(dict_of_neighbours[first_solution[-2]][position][1]    )
         - 1_00_00
     )
     return first_solution, distance_of_first_solution
def       UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_    ):



     """simple docstring"""






     snake_case					=    []

     for n in solution[1:-1]:
          snake_case					=    solution.index(UpperCamelCase_    )
          for kn in solution[1:-1]:
               snake_case					=    solution.index(UpperCamelCase_    )
               if n == kn:
                    continue

               snake_case					=    copy.deepcopy(UpperCamelCase_    )
               snake_case					=    kn
               snake_case					=    n

               snake_case					=    0

               for k in _tmp[:-1]:
                    snake_case					=    _tmp[_tmp.index(UpperCamelCase_    ) + 1]
                    for i in dict_of_neighbours[k]:
                         if i[0] == next_node:
                              snake_case					=    distance + int(i[1]    )
               _tmp.append(UpperCamelCase_    )

               if _tmp not in neighborhood_of_solution:
                    neighborhood_of_solution.append(_tmp    )

     snake_case					=    len(neighborhood_of_solution[0]    ) - 1

     neighborhood_of_solution.sort(key=lambda UpperCamelCase_    : x[index_of_last_item_in_the_list]    )
     return neighborhood_of_solution
def       UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_    ):



     """simple docstring"""






     snake_case					=    1
     snake_case					=    first_solution
     snake_case					=    []
     snake_case					=    distance_of_first_solution
     snake_case					=    solution

     while count <= iters:
          snake_case					=    find_neighborhood(UpperCamelCase_							,UpperCamelCase_    )
          snake_case					=    0
          snake_case					=    neighborhood[index_of_best_solution]
          snake_case					=    len(UpperCamelCase_    ) - 1

          snake_case					=    False
          while not found:
               snake_case					=    0
               while i < len(UpperCamelCase_    ):
                    if best_solution[i] != solution[i]:
                         snake_case					=    best_solution[i]
                         snake_case					=    solution[i]
                         break
                    snake_case					=    i + 1

               if [first_exchange_node, second_exchange_node] not in tabu_list and [
                   second_exchange_node,
                   first_exchange_node,
               ] not in tabu_list:
                    tabu_list.append([first_exchange_node, second_exchange_node]    )
                    snake_case					=    True
                    snake_case					=    best_solution[:-1]
                    snake_case					=    neighborhood[index_of_best_solution][best_cost_index]
                    if cost < best_cost:
                         snake_case					=    cost
                         snake_case					=    solution
               else:
                    snake_case					=    index_of_best_solution + 1
                    snake_case					=    neighborhood[index_of_best_solution]

          if len(UpperCamelCase_    ) >= size:
               tabu_list.pop(0    )

          snake_case					=    count + 1

     return best_solution_ever, best_cost
def       UpperCAmelCase__						(UpperCamelCase_=None    ):



     """simple docstring"""






     snake_case					=    generate_neighbours(args.File    )

     snake_case			,  snake_case					=    generate_first_solution(
         args.File							,UpperCamelCase_    )

     snake_case			,  snake_case					=    tabu_search(
         UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_							,args.Iterations							,args.Size							,)

     print(F'''Best solution: {best_sol}, with total distance: {best_cost}.'''    )


if __name__ == "__main__":
 _SCREAMING_SNAKE_CASE      :			Optional[Any]      				=   argparse.ArgumentParser(description="Tabu Search")
 parser.add_argument(
     "-f",
     "--File",
     type=str,
     help="Path to the file containing the data",
     required=True,
 )
 parser.add_argument(
     "-i",
     "--Iterations",
     type=int,
     help="How many iterations the algorithm should perform",
     required=True,
 )
 parser.add_argument(
     "-s", "--Size", type=int, help="Size of the tabu list", required=True
 )

 # Pass the arguments to main method
 main(parser.parse_args())




