# coding=utf-8
# Copyright 2018 The Google AI Language Team Authors,
# The HuggingFace Inc. team, and The XTREME Benchmark Authors.
# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Fine-tuning models for NER and POS tagging."""

from __future__ import absolute_import, division, print_function

import argparse
import glob
import logging
import os
import random
from dataclasses import dataclass, field
from typing import Optional
import json

import numpy as np
import scipy
import torch
from seqeval.metrics import precision_score, recall_score, f1_score
from tensorboardX import SummaryWriter
from torch.nn import CrossEntropyLoss
from torch.utils.data import DataLoader, TensorDataset
from torch.utils.data import RandomSampler, SequentialSampler
from torch.utils.data.distributed import DistributedSampler
from tqdm import tqdm, trange
from utils_tag import convert_examples_to_features
from utils_tag import get_labels
from utils_tag import read_examples_from_file
# import lang2vec.lang2vec as l2v
from scipy.spatial import distance

from transformers import (
  AdamW,
  get_linear_schedule_with_warmup,
  WEIGHTS_NAME,
  AutoConfig,
  AutoModelForTokenClassification,
  AutoTokenizer,
  HfArgumentParser,
  MultiLingAdapterArguments,
  AdapterConfig,
  AdapterType,
)
#from xlm import XLMForTokenClassification

DEFAULT_LANGUAGES = {
  'mr': 'hi',
  'bn': 'hi',
  'ta': 'ta',
  'fo': 'fo',
  'no': 'da',
  'da': 'da',
  'be': 'be',
  'uk': 'uk',
  'bg': 'bg'
}
logger = logging.getLogger(__name__)

def set_seed(args):
  random.seed(args.seed)
  np.random.seed(args.seed)
  torch.manual_seed(args.seed)
  logger.info(f'Seed = {args.seed}')
  if args.n_gpu > 0:
    torch.cuda.manual_seed_all(args.seed)


def train(args, train_dataset, model, tokenizer, labels, pad_token_label_id, lang_adapter_names, task_name, lang2id=None):
  """Train the model."""
  if args.local_rank in [-1, 0]:
    tb_writer = SummaryWriter()

  args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
  print(f'Local Rank = {args.local_rank}')
  print(len(train_dataset))
  train_sampler = RandomSampler(train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset)
  train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.train_batch_size)

  if args.max_steps > 0:
    t_total = args.max_steps
    args.num_train_epochs = args.max_steps // (len(train_dataloader) // args.gradient_accumulation_steps) + 1
  else:
    t_total = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs

  # Prepare optimizer and schedule (linear warmup and decay)
  no_decay = ["bias", "LayerNorm.weight"]
  optimizer_grouped_parameters = [
    {"params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
     "weight_decay": args.weight_decay},
    {"params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], "weight_decay": 0.0}
  ]
  optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
  logging.info([n for (n, p) in model.named_parameters() if p.requires_grad])
  scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=t_total)
  if args.fp16:
    try:
      from apex import amp
    except ImportError:
      raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
    model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level)

  # multi-gpu training (should be after apex fp16 initialization)
  if args.n_gpu > 1:
    model = torch.nn.DataParallel(model)

  # Distributed training (should be after apex fp16 initialization)
  if args.local_rank != -1:
    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank],
                              output_device=args.local_rank,
                              find_unused_parameters=True)

  # Train!
  logger.info("***** Running training *****")
  logger.info("  Num examples = %d", len(train_dataset))
  logger.info("  Num Epochs = %d", args.num_train_epochs)
  logger.info("  Instantaneous batch size per GPU = %d", args.per_gpu_train_batch_size)
  logger.info("  Total train batch size (w. parallel, distributed & accumulation) = %d",
        args.train_batch_size * args.gradient_accumulation_steps * (
          torch.distributed.get_world_size() if args.local_rank != -1 else 1))
  logger.info("  Gradient Accumulation steps = %d", args.gradient_accumulation_steps)
  logger.info("  Total optimization steps = %d", t_total)

  best_score = 0.0
  best_checkpoint = None
  patience = 0
  global_step = 0
  tr_loss, logging_loss = 0.0, 0.0
  model.zero_grad()
  train_iterator = trange(int(args.num_train_epochs), desc="Epoch", disable=args.local_rank not in [-1, 0])
  set_seed(args) # Add here for reproductibility (even between python 2 and 3)

  cur_epoch = 0
  for _ in train_iterator:
    epoch_iterator = tqdm(train_dataloader, desc="Iteration", disable=args.local_rank not in [-1, 0])
    cur_epoch += 1
    for step, batch in enumerate(epoch_iterator):
      batch = tuple(t.to(args.device) for t in batch if t is not None)
      inputs = {"input_ids": batch[0],
            "attention_mask": batch[1],
            "labels": batch[3]}

      if args.model_type != "distilbert":
        # XLM and RoBERTa don"t use segment_ids
        inputs["token_type_ids"] = batch[2] if args.model_type in ["bert", "xlnet"] else None

      if args.model_type == "xlm":
        inputs["langs"] = batch[4]

      outputs = model(**inputs)
      loss = outputs[0]

      if args.n_gpu > 1:
        # mean() to average on multi-gpu parallel training
        loss = loss.mean()
      if args.gradient_accumulation_steps > 1:
        loss = loss / args.gradient_accumulation_steps

      if args.fp16:
        with amp.scale_loss(loss, optimizer) as scaled_loss:
          scaled_loss.backward()
      else:
        loss.backward()
      tr_loss += loss.item()
      if (step + 1) % args.gradient_accumulation_steps == 0:
        if args.fp16:
          torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
        else:
          torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)

        scheduler.step()  # Update learning rate schedule
        optimizer.step()
        model.zero_grad()
        global_step += 1

        if args.local_rank in [-1, 0] and args.logging_steps > 0 and global_step % args.logging_steps == 0:
          # Log metrics
          if args.local_rank == -1 and args.evaluate_during_training:
            # Only evaluate on single GPU otherwise metrics may not average well
            results, _ = evaluate(args, model, tokenizer, labels, pad_token_label_id, mode="dev", lang=args.train_langs, lang2id=lang2id, lang_adapter_names=lang_adapter_names, task_name=task_name)
            for key, value in results.items():
              tb_writer.add_scalar("eval_{}".format(key), value, global_step)
          tb_writer.add_scalar("lr", scheduler.get_lr()[0], global_step)
          tb_writer.add_scalar("loss", (tr_loss - logging_loss) / args.logging_steps, global_step)
          logging_loss = tr_loss

        if args.local_rank in [-1, 0] and args.save_steps > 0 and global_step % args.save_steps == 0:
          if args.save_only_best_checkpoint:
            result, _ = evaluate(args, model, tokenizer, labels, pad_token_label_id, mode="dev", prefix=global_step, lang=args.train_langs, lang2id=lang2id, lang_adapter_names=lang_adapter_names, task_name=task_name)
            if result["f1"] > best_score:
              logger.info("result['f1']={} > best_score={}".format(result["f1"], best_score))
              best_score = result["f1"]
              # Save the best model checkpoint
              output_dir = os.path.join(args.output_dir, "checkpoint-best")
              best_checkpoint = output_dir
              if not os.path.exists(output_dir):
                os.makedirs(output_dir)
              # Take care of distributed/parallel training
              model_to_save = model.module if hasattr(model, "module") else model
              if args.do_save_adapters:
                model_to_save.save_all_adapters(output_dir)
              if args.do_save_adapter_fusions:
                model_to_save.save_all_adapter_fusions(output_dir)
              if args.do_save_full_model:
                model_to_save.save_pretrained(output_dir)
              torch.save(args, os.path.join(output_dir, "training_args.bin"))
              logger.info("Saving the best model checkpoint to %s", output_dir)
              logger.info("Reset patience to 0")
              patience = 0
            else:
              patience += 1
              logger.info("Hit patience={}".format(patience))
              if args.eval_patience > 0 and patience > args.eval_patience:
                logger.info("early stop! patience={}".format(patience))
                epoch_iterator.close()
                train_iterator.close()
                if args.local_rank in [-1, 0]:
                  tb_writer.close()
                return global_step, tr_loss / global_step
          else:
            # Save model checkpoint
            output_dir = os.path.join(args.output_dir, "checkpoint-{}".format(global_step))
            if not os.path.exists(output_dir):
              os.makedirs(output_dir)
            # Take care of distributed/parallel training
            model_to_save = model.module if hasattr(model, "module") else model
            if args.do_save_adapters:
              model_to_save.save_all_adapters(output_dir)
            if args.do_save_adapter_fusions:
              model_to_save.save_all_adapter_fusions(output_dir)
            if args.do_save_full_model:
              model_to_save.save_pretrained(output_dir)
            torch.save(args, os.path.join(output_dir, "training_args.bin"))
            logger.info("Saving model checkpoint to %s", output_dir)

      if args.max_steps > 0 and global_step > args.max_steps:
        epoch_iterator.close()
        break
    if args.max_steps > 0 and global_step > args.max_steps:
      train_iterator.close()
      break

  if args.local_rank in [-1, 0]:
    tb_writer.close()

  return global_step, tr_loss / global_step

def calc_weight_multi(args, model, batch, lang_adapter_names, task_name, adapter_weights, step=10, lang=None):
  inputs = {"input_ids": batch[0],
        "attention_mask": batch[1],
        "return_sequence_out": True,
        "labels": batch[3]}
  # logger.info(f'Language Adapters are {lang_adapter_names}')
  adapter_weights = [torch.FloatTensor([0.5 for _ in range(len(lang_adapter_names))]).to(args.device) for _ in range(13)]
  if args.lang_to_vec:
    logger.info(lang)
    logger.info(lang_adapter_names)
    adapter_weights = calc_l2v_weights(lang, lang_adapter_names, args.en_weight)
    logger.info(adapter_weights)
  for step_no in range(step):
    for w in adapter_weights: w.requires_grad = True
    if args.lang_to_vec and step_no == 0:
      normed_adapter_weights = adapter_weights
    else:
      normed_adapter_weights = [torch.nn.functional.softmax(w) for w in adapter_weights]
    # logger.info(f'Initial Adapter Weights = {normed_adapter_weights}')
    model.set_active_adapters([lang_adapter_names, [task_name]])
    inputs["adapter_names"] = [lang_adapter_names, [task_name]]

    inputs["adapter_weights"] = normed_adapter_weights
    outputs = model(**inputs)

    loss, logits, orig_sequence_output = outputs[:3]
    kept_logits = outputs[-1]
    entropy = torch.nn.functional.softmax(kept_logits, dim=1)*torch.nn.functional.log_softmax(kept_logits, dim=1)
    entropy = -entropy.sum() / kept_logits.size(0)
    grads = torch.autograd.grad(entropy, adapter_weights)
    #print(adapter_weights)
    #print(grads)
    #print(grads)
    for i, w in enumerate(adapter_weights):
      adapter_weights[i] = adapter_weights[i].data - 10*grads[i].data


  normed_adapter_weights = [torch.nn.functional.softmax(w) for w in adapter_weights]
  #print(normed_adapter_weights)
  # logger.info(f'Final Adapter Weights = {normed_adapter_weights}')
  return normed_adapter_weights

def jaccard_sim(vec1, vec2):
    intersection = 0
    union = 0
    for i in range(len(vec1)):
        if vec1[i] == '--' or vec2[i] == '--':
            continue
        if vec1[i] == 1 or vec2[i] == 1:
            union += 1
        if vec1[i] == 1 and vec2[i] == 1:
            intersection += 1
    return intersection/union

def get_sim(lang1, lang2):
  features = l2v.get_features(f'{DEFAULT_LANGUAGES[lang1]} {lang2}', 'learned')
  similarity = 1 - distance.cosine(features[DEFAULT_LANGUAGES[lang1]], features[lang2])
  return similarity

def get_syntax_sim(lang1, lang2):
  features = l2v.get_features(f'{lang1} {lang2}', "syntax_wals|syntax_sswl|syntax_ethnologue")
  similarity = jaccard_sim(features[lang1], features[lang2])
  return similarity

def calc_l2v_weights(args, lang, lang_adapter_names):
  adapter_weight = []
  for adapter_lang in lang_adapter_names:
    if args.en_weight is not None and adapter_lang == 'en':
      continue
    if args.lang_to_vec == 'learned':
      adapter_weight.append(get_sim(lang, adapter_lang))
    elif args.lang_to_vec == 'syntax':
      adapter_weight.append(get_syntax_sim(lang, adapter_lang))
    else:
      logger.info('INVALID FEATURE TYPE')
      exit()
  logger.info(adapter_weight)
  adapter_weight = torch.FloatTensor(adapter_weight)
  adapter_weight = torch.nn.functional.softmax(adapter_weight/args.temperature).tolist()
  if args.en_weight is not None:
    adapter_weight = [(1 - args.en_weight)*aw for aw in adapter_weight]
    en_index = lang_adapter_names.index('en')
    adapter_weight.insert(en_index, args.en_weight)
  return adapter_weight

def scaled_input(emb, batch_size=16, num_batch=1, baseline=None, start_i=None, end_i=None):
    # shape of emb: (num_head, seq_len, seq_len)
    if baseline is None:
        baseline = torch.zeros_like(emb)   

    num_points = batch_size * num_batch
    scale = 1.0 / num_points
    if start_i is None:
        step = (emb.unsqueeze(0) - baseline.unsqueeze(0)) * scale
        res = torch.cat([torch.add(baseline.unsqueeze(0), step*i) for i in range(num_points)], dim=0)
        return res, step[0]
    else:
        step = (emb - baseline) * scale
        start_emb = torch.add(baseline, step*start_i)
        end_emb = torch.add(baseline, step*end_i)
        step_new = (end_emb.unsqueeze(0) - start_emb.unsqueeze(0)) * scale
        res = torch.cat([torch.add(start_emb.unsqueeze(0), step_new*i) for i in range(num_points)], dim=0)
        return res, step_new[0]

#Changed the default of calc_weight_step to 0
def evaluate(args, model, tokenizer, labels, pad_token_label_id, mode, prefix="", lang="en", lang2id=None, print_result=True, adapter_weight=None, lang_adapter_names=None, task_name=None, calc_weight_step=0):
  eval_dataset = load_and_cache_examples(args, tokenizer, labels, pad_token_label_id, mode=mode, lang=lang, lang2id=lang2id)
  args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)
  # Note that DistributedSampler samples randomly
  if args.get_attr:
    eval_sampler = RandomSampler(eval_dataset) if args.local_rank == -1 else DistributedSampler(eval_dataset)
  else:
    eval_sampler = SequentialSampler(eval_dataset) if args.local_rank == -1 else DistributedSampler(eval_dataset)
  eval_dataloader = DataLoader(eval_dataset, sampler=eval_sampler, batch_size=args.eval_batch_size)

  # multi-gpu evaluate
  if args.n_gpu > 1:
    model = torch.nn.DataParallel(model)
  # Eval!
  logger.info("***** Running evaluation %s in %s *****" % (prefix, lang))
  logger.info("  Num examples = %d", len(eval_dataset))
  logger.info("  Batch size = %d", args.eval_batch_size)
  eval_loss = 0.0
  nb_eval_steps = 0
  preds = None
  out_label_ids = None
  model.eval()
  counter = 0
  
  head_importances = None
  all_head_importances = None
  for batch in tqdm(eval_dataloader, desc="Evaluating"):
    counter += 1 
    logger.info(f'Batch number = {counter}')
    batch = tuple(t.to(args.device) for t in batch)
    if calc_weight_step > 0:
      adapter_weight = calc_weight_multi(args, model, batch, lang_adapter_names, task_name, adapter_weight, calc_weight_step, lang=lang)
    if args.get_attr:
      inputs = {"input_ids": batch[0],
            "attention_mask": batch[1],
            "labels": batch[3],
            "adapter_weights": adapter_weight}
      if args.model_type != "distilbert":
        # XLM and RoBERTa don"t use segment_ids
        inputs["token_type_ids"] = batch[2] if args.model_type in ["bert", "xlnet"] else None
      if args.model_type == 'xlm':
        inputs["langs"] = batch[4]
      inputs["output_attentions"] = True
      outputs = model(**inputs)
      tmp_eval_loss, logits, attentions, kept_labels, kl_logits  = outputs

      attr_all = []
      res_attr = []
      
      input_len = int(inputs["attention_mask"][0].sum())
      example_head_importances = None
      #Remove the batch_size dim since batch_size=1
      logits = logits[0]
      for tar_layer in range(12):
        att = attentions[tar_layer][0]
        pred_labels = torch.argmax(logits, dim=-1)

        scale_att, step = scaled_input(att.data)
        scale_att.requires_grad_(True)

        attr_all = None
        prob_all = None
        for j_batch in range(1):
            one_batch_att = scale_att[j_batch*16:(j_batch+1)*16]
            _, grad = model(input_ids=inputs['input_ids'], token_type_ids=inputs['token_type_ids'], attention_mask=inputs['attention_mask'], labels=inputs['labels'], tar_layer=tar_layer, tmp_score=one_batch_att, pred_labels=pred_labels)
            grad = grad.sum(dim=0) 
            attr_all = grad if attr_all is None else torch.add(attr_all, grad)
            # prob_all = tar_prob if prob_all is None else torch.cat([prob_all, tar_prob])
        
        attr_all = attr_all[:,0:input_len,0:input_len] * step[:,0:input_len,0:input_len]
        if example_head_importances is None:
          example_head_importances = torch.amax(attr_all, dim=(1,2)).unsqueeze(0)
        else:
          tmp = torch.amax(attr_all, dim=(1,2))
          tmp = tmp.unsqueeze(0)
          example_head_importances = torch.cat((example_head_importances, tmp), dim=0)
        # att = att[:,0:input_len,0:input_len]
        res_attr.append(attr_all.data)
      # logger.info(f'Example Head Importances = {example_head_importances}')
      all_head_importances = example_head_importances.unsqueeze(0) if all_head_importances is None else torch.cat((all_head_importances, example_head_importances.unsqueeze(0)), dim=0)
      head_importances = example_head_importances if head_importances is None else torch.add(head_importances, example_head_importances)
      if counter == 100:
        break
      continue
    with torch.no_grad():
      inputs = {"input_ids": batch[0],
            "attention_mask": batch[1],
            "labels": batch[3],
            "adapter_weights": adapter_weight}
      # logger.info(f'Labels = {batch[3]}')
      if args.model_type != "distilbert":
        # XLM and RoBERTa don"t use segment_ids
        inputs["token_type_ids"] = batch[2] if args.model_type in ["bert", "xlnet"] else None
      if args.model_type == 'xlm':
        inputs["langs"] = batch[4]


      outputs = model(**inputs)
      tmp_eval_loss, logits = outputs[:2]

      if args.n_gpu > 1:
        # mean() to average on multi-gpu parallel evaluating
        tmp_eval_loss = tmp_eval_loss.mean()

      eval_loss += tmp_eval_loss.item()
    nb_eval_steps += 1
    if preds is None:
      preds = logits.detach().cpu().numpy()
      out_label_ids = inputs["labels"].detach().cpu().numpy()
    else:
      preds = np.append(preds, logits.detach().cpu().numpy(), axis=0)
      out_label_ids = np.append(out_label_ids, inputs["labels"].detach().cpu().numpy(), axis=0)

  if args.get_attr:
    head_importances = head_importances/counter
    logger.info(f'Head Importances = {head_importances}')
    torch.save(head_importances, os.path.join(args.output_dir,f'{mode}_{lang}_s{args.seed}_importances_100.pt'))
    torch.save(all_head_importances, os.path.join(args.output_dir,f'{mode}_{lang}_s{args.seed}_all_importances_100.pt'))
    return None, None

  if nb_eval_steps == 0:
    results = {k: 0 for k in ["loss", "precision", "recall", "f1"]}
  else:
    eval_loss = eval_loss / nb_eval_steps
    preds = np.argmax(preds, axis=2)

    label_map = {i: label for i, label in enumerate(labels)}

    out_label_list = [[] for _ in range(out_label_ids.shape[0])]
    preds_list = [[] for _ in range(out_label_ids.shape[0])]

    for i in range(out_label_ids.shape[0]):
      for j in range(out_label_ids.shape[1]):
        if out_label_ids[i, j] != pad_token_label_id:
          out_label_list[i].append(label_map[out_label_ids[i][j]])
          preds_list[i].append(label_map[preds[i][j]])

    results = {
      "loss": eval_loss,
      "precision": precision_score(out_label_list, preds_list),
      "recall": recall_score(out_label_list, preds_list),
      "f1": f1_score(out_label_list, preds_list)
    }

  if print_result:
    logger.info("***** Evaluation result %s in %s *****" % (prefix, lang))
    for key in sorted(results.keys()):
      logger.info("  %s = %s", key, str(results[key]))
  return results, preds_list


def load_and_cache_examples(args, tokenizer, labels, pad_token_label_id, mode, lang, lang2id=None, few_shot=-1):
  # Make sure only the first process in distributed training process
  # the dataset, and the others will use the cache
  if args.local_rank not in [-1, 0] and not evaluate:
    torch.distributed.barrier()

  # Load data features from cache or dataset file
  bpe_dropout = args.bpe_dropout
  if mode != 'train': bpe_dropout = 0
  if bpe_dropout > 0:
    cached_features_file = os.path.join(args.data_dir, "cached_{}_{}_{}_{}_drop{}".format(mode, lang,
      list(filter(None, args.model_name_or_path.split("/"))).pop(),
      str(args.max_seq_length), bpe_dropout))
  else:
    cached_features_file = os.path.join(args.data_dir, "cached_{}_{}_{}_{}".format(mode, lang,
      list(filter(None, args.model_name_or_path.split("/"))).pop(),
      str(args.max_seq_length)))
  if os.path.exists(cached_features_file) and not args.overwrite_cache:
    logger.info("Loading features from cached file %s", cached_features_file)
    features = torch.load(cached_features_file)
  else:
    langs = lang.split(',')
    logger.info("all languages = {}".format(lang))
    features = []
    for lg in langs:
      data_file = os.path.join(args.data_dir, lg, "{}.{}".format(mode, args.model_name_or_path))
      logger.info("Creating features from dataset file at {} in language {}".format(data_file, lg))
      examples = read_examples_from_file(data_file, lg, lang2id)
      print(examples)
      features_lg = convert_examples_to_features(examples, labels, args.max_seq_length, tokenizer,
                          cls_token_at_end=bool(args.model_type in ["xlnet"]),
                          cls_token=tokenizer.cls_token,
                          cls_token_segment_id=2 if args.model_type in ["xlnet"] else 0,
                          sep_token=tokenizer.sep_token,
                          sep_token_extra=bool(args.model_type in ["roberta", "xlmr"]),
                          pad_on_left=bool(args.model_type in ["xlnet"]),
                          pad_token=tokenizer.convert_tokens_to_ids([tokenizer.pad_token])[0],
                          pad_token_segment_id=4 if args.model_type in ["xlnet"] else 0,
                          pad_token_label_id=pad_token_label_id,
                          lang=lg,
                          bpe_dropout=bpe_dropout,
                          )
      features.extend(features_lg)
    if args.local_rank in [-1, 0]:
      logger.info("Saving features into cached file {}, len(features)={}".format(cached_features_file, len(features)))
      torch.save(features, cached_features_file)

  # Make sure only the first process in distributed training process
  # the dataset, and the others will use the cache
  if args.local_rank == 0 and not evaluate:
    torch.distributed.barrier()

  if few_shot > 0 and mode == 'train':
    logger.info("Original no. of examples = {}".format(len(features)))
    features = features[: few_shot]
    logger.info('Using few-shot learning on {} examples'.format(len(features)))

  # Convert to Tensors and build dataset
  all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)
  all_input_mask = torch.tensor([f.input_mask for f in features], dtype=torch.long)
  all_segment_ids = torch.tensor([f.segment_ids for f in features], dtype=torch.long)
  all_label_ids = torch.tensor([f.label_ids for f in features], dtype=torch.long)
  if args.model_type == 'xlm' and features[0].langs is not None:
    all_langs = torch.tensor([f.langs for f in features], dtype=torch.long)
    logger.info('all_langs[0] = {}'.format(all_langs[0]))
    dataset = TensorDataset(all_input_ids, all_input_mask, all_segment_ids, all_label_ids, all_langs)
  else:
    dataset = TensorDataset(all_input_ids, all_input_mask, all_segment_ids, all_label_ids)
  return dataset


@dataclass
class ModelArguments:
    """
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
    """

    model_name_or_path: str = field(
        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
    )
    model_type: Optional[str] = field(
        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
    )
    config_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
    )
    tokenizer_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
    )
    cache_dir: Optional[str] = field(
        default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}
    )
    labels: str = field(
        default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}
    )
    data_dir: str = field(
        default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}
    )
    output_dir: str = field(
        default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}
    )
    max_seq_length: Optional[int] = field(
        default=128, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}
    )
    do_train: Optional[bool] = field(default=False )
    do_eval: Optional[bool] = field(default=False )
    do_predict: Optional[bool] = field(default=False )
    do_adapter_predict: Optional[bool] = field(default=False )
    do_predict_dev: Optional[bool] = field(default=False )
    do_predict_train: Optional[bool] = field(default=False )
    init_checkpoint: Optional[str] = field(default=None )
    evaluate_during_training: Optional[bool] = field(default=False )
    do_lower_case: Optional[bool] = field(default=False )
    few_shot: Optional[int] = field(default=-1 )
    per_gpu_train_batch_size: Optional[int] = field(default=8)
    per_gpu_eval_batch_size: Optional[int] = field(default=8)
    gradient_accumulation_steps: Optional[int] = field(default=1)
    learning_rate: Optional[float] = field(default=5e-5)
    weight_decay: Optional[float] = field(default=0.0)
    adam_epsilon: Optional[float] = field(default=1e-8)
    max_grad_norm: Optional[float] = field(default=1.0)
    num_train_epochs: Optional[float] = field(default=3.0)
    max_steps: Optional[int] = field(default=-1)
    save_steps: Optional[int] = field(default=-1)
    warmup_steps: Optional[int] = field(default=0)
    logging_steps: Optional[int] = field(default=50)
    save_only_best_checkpoint: Optional[bool] = field(default=False)
    eval_all_checkpoints: Optional[bool] = field(default=False)
    no_cuda: Optional[bool] = field(default=False)
    overwrite_output_dir: Optional[bool] = field(default=False)
    overwrite_cache: Optional[bool] = field(default=False)
    seed: Optional[int] = field(default=42)
    fp16: Optional[bool] = field(default=False)
    fp16_opt_level: Optional[str] = field(default="O1")
    local_rank: Optional[int] = field(default=-1)
    server_ip: Optional[str] = field(default="")
    server_port: Optional[str] = field(default="")
    predict_langs: Optional[str] = field(default="en")
    train_langs: Optional[str] = field(default="en")
    log_file: Optional[str] = field(default=None)
    eval_patience: Optional[int] = field(default=-1)
    bpe_dropout: Optional[float] = field(default=0)
    do_save_adapter_fusions: Optional[bool] = field(default=False)
    task_name: Optional[str] = field(default="ner")

    predict_task_adapter: Optional[str] = field(default=None)
    predict_lang_adapter: Optional[str] = field(default=None)
    test_adapter: Optional[bool] = field(default=False)

    adapter_weight: Optional[str] = field(default=None)
    lang_to_vec: Optional[str] = field(default=None)

    calc_weight_step: Optional[int] = field(default=0)
    predict_save_prefix: Optional[str] = field(default=None)
    en_weight: Optional[float] = field(default=None)
    temperature: Optional[float] = field(default=1.0)

    get_attr: Optional[bool] = field(default=False)
    topk: Optional[int] = field(default=1)

    task: Optional[str] = field(default='udpos')

def setup_adapter(args, adapter_args, model, train_adapter=True, load_adapter=None, load_lang_adapter=None):
  task_name = args.task_name or "ner"
  # check if adapter already exists, otherwise add it
  if task_name not in model.config.adapters.adapter_list(AdapterType.text_task):
      logging.info("Trying to decide if add adapter")
      # resolve the adapter config
      adapter_config = AdapterConfig.load(
          adapter_args.adapter_config,
          non_linearity=adapter_args.adapter_non_linearity,
          reduction_factor=adapter_args.adapter_reduction_factor,
      )
      # load a pre-trained from Hub if specified
      if adapter_args.load_adapter or load_adapter:
          logging.info("loading task adapter")
          model.load_adapter(
              adapter_args.load_adapter if load_adapter is None else load_adapter,
              AdapterType.text_task,
              config=adapter_config,
              load_as=task_name,
          )
      # otherwise, add a fresh adapter
      else:
          logging.info("Adding task adapter")
          model.add_adapter(task_name, AdapterType.text_task, config=adapter_config)
  # optionally load a pre-trained language adapter
  if adapter_args.load_lang_adapter or load_lang_adapter:

      if load_lang_adapter is None:
          # load a set of language adapters
          logging.info("loading lang adpater {}".format(adapter_args.load_lang_adapter))
          # resolve the language adapter config
          lang_adapter_config = AdapterConfig.load(
              adapter_args.lang_adapter_config,
              non_linearity=adapter_args.lang_adapter_non_linearity,
              reduction_factor=adapter_args.lang_adapter_reduction_factor,
          )
          # load the language adapter from Hub
          # if adapter_args.language == 'topk':
          #   assert len(args.predict_langs.split(',')) == 1
          #   filename = f'scripts/{args.task}/en/{args.predict_langs}.json'
          #   logger.info(f'Loading Adapter Languages from {filename}')
          #   languages = []
          #   with open(filename) as f:
          #     for i,line in enumerate(f):
          #       if i == args.topk:
          #         break
          #       line = json.loads(line)
          #       languages.append(line['adapter'].strip())
          #   adapter_names = [f'{lang}/wiki@ukp' for lang in languages]
          # else:
          #   languages = adapter_args.language.split(",")
          #   adapter_names = adapter_args.load_lang_adapter.split(",")
          # logger.info(f'Adapter Languages : {languages}, Length : {len(languages)}')
          # logger.info(f'Adapter Names {adapter_names}, Length : {len(adapter_names)}')
          # assert len(languages) == len(adapter_names)
          # lang_adapter_names = []
          # for language, adapter_name in zip(languages, adapter_names):
          #     logger.info(f'Language = {language}')
          #     logger.info(f'Adapter Name = {adapter_name}')
          #     lang_adapter_name = model.load_adapter(
          #         adapter_name,
          #         AdapterType.text_lang,
          #         config=lang_adapter_config,
          #         load_as=language,
          #     )
          #     lang_adapter_names.append(lang_adapter_name)
      else:
          logging.info("loading lang adpater {}".format(load_lang_adapter))
          # resolve the language adapter config
          lang_adapter_config = AdapterConfig.load(
              adapter_args.lang_adapter_config,
              non_linearity=adapter_args.lang_adapter_non_linearity,
              reduction_factor=adapter_args.lang_adapter_reduction_factor,
          )
          # load the language adapter from Hub
          # lang_adapter_name = model.load_adapter(
          #     load_lang_adapter,
          #     AdapterType.text_lang,
          #     config=lang_adapter_config,
          #     load_as="lang",
          # )
          # lang_adapter_names = [lang_adapter_name]
  else:
      lang_adapter_name = None
      lang_adapter_names = []
  # Freeze all model weights except of those of this adapter
  model.train_adapter([task_name])

  # Set the adapters to be used in every forward pass
  if lang_adapter_name:
      model.set_active_adapters([lang_adapter_names, [task_name]])
  else:
      model.set_active_adapters([task_name])

  return model, lang_adapter_names, task_name

def load_model(args, num_labels):
  logger.info('Loading pretrained model and tokenizer')
  config = AutoConfig.from_pretrained(
    args.config_name if args.config_name else args.model_name_or_path,
    num_labels=num_labels,
    cache_dir=args.cache_dir,
  )
  args.model_type = config.model_type
  tokenizer = AutoTokenizer.from_pretrained(
    args.tokenizer_name if args.tokenizer_name else args.model_name_or_path,
    do_lower_case=args.do_lower_case,
    cache_dir=args.cache_dir,
    use_fast=False,
  )
  if args.init_checkpoint:
    logger.info("loading from init_checkpoint={}".format(args.init_checkpoint))
    model = AutoModelForTokenClassification.from_pretrained(
        args.init_checkpoint,
        config=config,
        cache_dir=args.cache_dir,
    )
  else:
    logger.info("loading from existing model {}".format(args.model_name_or_path))
    model = AutoModelForTokenClassification.from_pretrained(
        args.model_name_or_path,
        from_tf=bool(".ckpt" in args.model_name_or_path),
        config=config,
        cache_dir=args.cache_dir,
    )
  lang2id = config.lang2id if args.model_type == "xlm" else None
  logger.info("Using lang2id = {}".format(lang2id))

  return model, tokenizer, lang2id

def predict_and_save(args, adapter_args, model, tokenizer, labels, lang2id, pad_token_label_id, lang_adapter_names, task_name, split):
  output_test_results_file = os.path.join(args.output_dir, f"{split}_results.txt")
  with open(output_test_results_file, "a") as result_writer:
    for lang in args.predict_langs.split(','):
      #Check if language data exists
      if not os.path.exists(os.path.join(args.data_dir, lang, '{}.{}'.format(split, args.model_name_or_path))):
        logger.info("Language {}, split {} does not exist".format(lang, split))
        continue

      #Activate the required language adapter
      adapter_weight = None
    #   if not args.adapter_weight and not args.lang_to_vec:
    #     if (adapter_args.train_adapter or args.test_adapter) and not args.adapter_weight:
    #       if lang in lang_adapter_names:
    #         logger.info(f'Language adapter for {lang} found')
    #         logger.info("Set active language adapter to {}".format(lang))
    #         model.set_active_adapters([[lang], [task_name]])
    #       else:
    #         logger.info(f'Language adapter for {lang} not found, using {lang_adapter_names[0]} instead')
    #         logger.info("Set active language adapter to {}".format(lang_adapter_names[0]))
    #         model.set_active_adapters([[lang_adapter_names[0]], [task_name]])
    #   else:
    #     if args.adapter_weight == 'equal':
    #       adapter_weight = [1/len(lang_adapter_names) for _ in lang_adapter_names]
    #     elif args.adapter_weight == 'equal_en':
    #       assert 'en' in lang_adapter_names, 'English language adapter not included'
    #       adapter_weight = [(1-args.en_weight)/(len(lang_adapter_names)-1) for _ in lang_adapter_names]
    #       en_index = lang_adapter_names.index('en')
    #       adapter_weight[en_index] = args.en_weight
    #     elif args.lang_to_vec:
    #       if args.en_weight is not None:
    #         logger.info(lang_adapter_names)
    #         assert 'en' in lang_adapter_names, 'English language adapter not included'
    #       adapter_weight = calc_l2v_weights(args, lang, lang_adapter_names)
    #     elif args.adapter_weight == 'load':
    #       filename = f'weights/{args.task}/{lang}/weights_s{args.seed}'
    #       logger.info(f'Loading adapter weights from {filename}')
    #       with open(filename) as f:
    #         adapter_weight = json.loads(next(f))
    #     elif args.adapter_weight != "0" and args.adapter_weight is not None:
    #       adapter_weight = [float(w) for w in args.adapter_weight.split(",")]
      logger.info('Args Adapter Weight = {}'.format(args.adapter_weight))
      logger.info('Adapter Languages = {}'.format(lang_adapter_names))
      if adapter_weight is not None:
        logger.info("Adapter Weights = {}".format(adapter_weight))
        logger.info('Sum of Adapter Weights = {}'.format(sum(adapter_weight)))
        logger.info("Length of Adapter Weights = {}".format(len(adapter_weight))) 
    #   model.set_active_adapters([ lang_adapter_names, [task_name]])
      #Evaluate
      result, predictions = evaluate(args, model, tokenizer, labels, pad_token_label_id, mode=split, lang=lang, lang2id=lang2id, adapter_weight=adapter_weight, lang_adapter_names=lang_adapter_names, task_name=task_name, calc_weight_step=args.calc_weight_step)

      if args.get_attr:
        continue
      result_json = {}
      # Save results
      if args.predict_save_prefix is not None and args.predict_save_prefix:
        result_json['language'] = f'{args.predict_save_prefix}_{lang}'
      else:
        result_json['language'] = f'{lang}'
      
      result_json['seed'] = args.seed
      result_json['language_adapters'] = lang_adapter_names
      if args.adapter_weight:
        result_json['adapter_weights'] = args.adapter_weight
      
      for key in sorted(result.keys()):
        result_json[key] = result[key]
      
      result_writer.write(json.dumps(result_json) + '\n')
      # Save predictions
      if args.predict_save_prefix is not None and args.predict_save_prefix:
        output_test_predictions_file = os.path.join(args.output_dir, "{}_{}_{}_s{}_predictions.txt".format(split, args.predict_save_prefix, lang, args.seed))
      else:
        output_test_predictions_file = os.path.join(args.output_dir, "{}_{}_s{}_predictions.txt".format(split, lang, args.seed))
      infile = os.path.join(args.data_dir, lang, "{}.{}".format(split, args.model_name_or_path))
      idxfile = infile + '.idx'
      save_predictions(args, predictions, output_test_predictions_file, infile, idxfile)

def main():
  parser = argparse.ArgumentParser()

  parser = HfArgumentParser((ModelArguments, MultiLingAdapterArguments))
  args, adapter_args = parser.parse_args_into_dataclasses()


  if os.path.exists(args.output_dir) and os.listdir(
      args.output_dir) and args.do_train and not args.overwrite_output_dir:
    raise ValueError(
      "Output directory ({}) already exists and is not empty. Use --overwrite_output_dir to overcome.".format(
        args.output_dir))

  # Setup distant debugging if needed
  if args.server_ip and args.server_port:
    import ptvsd
    print("Waiting for debugger attach")
    ptvsd.enable_attach(address=(args.server_ip, args.server_port), redirect_output=True)
    ptvsd.wait_for_attach()

  # Setup CUDA, GPU & distributed training
  if args.local_rank == -1 or args.no_cuda:
    device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
    args.n_gpu = torch.cuda.device_count()
  else:
  # Initializes the distributed backend which sychronizes nodes/GPUs
    torch.cuda.set_device(args.local_rank)
    device = torch.device("cuda", args.local_rank)
    torch.distributed.init_process_group(backend="nccl")
    args.n_gpu = 1
  args.device = device

  # Setup logging
  logging.basicConfig(handlers = [logging.FileHandler(args.log_file), logging.StreamHandler()],
                      format = '%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                      datefmt = '%m/%d/%Y %H:%M:%S',
                      level = logging.INFO if args.local_rank in [-1, 0] else logging.WARN)
  logging.info("Input args: %r" % args)
  logger.warning("Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
           args.local_rank, device, args.n_gpu, bool(args.local_rank != -1), args.fp16)
  # Set seed
  set_seed(args)
  # Prepare NER/POS task
  labels = get_labels(args.labels)
  num_labels = len(labels)
  # Use cross entropy ignore index as padding label id
  # so that only real label ids contribute to the loss later
  pad_token_label_id = CrossEntropyLoss().ignore_index

  # Load pretrained model and tokenizer
  # Make sure only the first process in distributed training loads model/vocab
  if args.local_rank not in [-1, 0]:
    torch.distributed.barrier() 

  args.do_save_full_model= (not adapter_args.train_adapter)
  args.do_save_adapters=adapter_args.train_adapter
  if args.do_save_adapters:
      logging.info('save adapters')
      logging.info(adapter_args.train_adapter)
  if args.do_save_full_model:
      logging.info('save model')

  # Make sure only the first process in distributed training loads model/vocab
  if args.local_rank == 0:
    torch.distributed.barrier()

  logger.info("Training/evaluation parameters %s", args)
  # Training
  if args.do_train:
    model, tokenizer, lang2id = load_model(args, num_labels)
    if adapter_args.train_adapter:
      model, lang_adapter_names, task_name = setup_adapter(args, adapter_args, model)
      logger.info("lang adapter names: {}".format(" ".join(lang_adapter_names)))
    else:
      lang_adatper_names = []
      task_name = None
    
    model.to(args.device)

    train_dataset = load_and_cache_examples(args, tokenizer, labels, pad_token_label_id, mode="train", lang=args.train_langs, lang2id=lang2id, few_shot=args.few_shot)
    global_step, tr_loss = train(args, train_dataset, model, tokenizer, labels, pad_token_label_id, lang_adapter_names, task_name, lang2id)
    logger.info(" global_step = %s, average loss = %s", global_step, tr_loss)

  # Saving best-practices: if you use default names for the model,
  # you can reload it using from_pretrained()
  if args.do_train and (args.local_rank == -1 or torch.distributed.get_rank() == 0):
    # Create output directory if needed
    if not os.path.exists(args.output_dir) and args.local_rank in [-1, 0]:
      os.makedirs(args.output_dir)

    # Save model, configuration and tokenizer using `save_pretrained()`.
    # They can then be reloaded using `from_pretrained()`
    # Take care of distributed/parallel training
    logger.info("Saving model checkpoint to %s", args.output_dir)
    model_to_save = model.module if hasattr(model, "module") else model
    if args.do_save_adapters:
      logging.info("Save adapter")
      model_to_save.save_all_adapters(args.output_dir)
    if args.do_save_adapter_fusions:
      logging.info("Save adapter fusion")
      model_to_save.save_all_adapter_fusions(args.output_dir)
    if args.do_save_full_model:
      logging.info("Save full model")
      model_to_save.save_pretrained(args.output_dir)

    tokenizer.save_pretrained(args.output_dir)

    # Good practice: save your training arguments together with the model
    torch.save(args, os.path.join(args.output_dir, "training_args.bin"))

  # Initialization for evaluation
  results = {}
  if args.init_checkpoint:
    best_checkpoint = args.init_checkpoint
  elif os.path.exists(os.path.join(args.output_dir, 'checkpoint-best')):
    best_checkpoint = os.path.join(args.output_dir, 'checkpoint-best')
  else:
    best_checkpoint = args.output_dir

  # Evaluation
  #This evaluates only if the entire model is saved, something we are not doing
  if args.do_eval and args.local_rank in [-1, 0]:
    model, tokenizer, lang2id = load_model(args, num_labels)

    logger.info('Evaluating the model on dev set of training language(en)')
    load_adapter = (best_checkpoint + "/" + args.task_name) if args.predict_task_adapter is None else args.predict_task_adapter
    # load_adapter = 'output/panx/bert-base-multilingual-cased-LR1e-4-epoch100-MaxLen128-TrainLangen_en_s0/checkpoint-best/ner/'
    logger.info(f'Task Adapter will be loaded from this path {load_adapter}')
    model.model_name = args.model_name_or_path
    model, lang_adapter_names, task_name = setup_adapter(args, adapter_args, model, load_adapter=load_adapter)
    model.to(args.device)
    
    result, _ = evaluate(args, model, tokenizer, labels, pad_token_label_id, mode="dev", prefix='debugging', lang=args.train_langs, lang2id=lang2id, lang_adapter_names=lang_adapter_names, task_name=task_name, calc_weight_step=args.calc_weight_step)
    results.update(result)
    # for checkpoint in checkpoints:
      # global_step = checkpoint.split("-")[-1] if len(checkpoints) > 1 else ""
      # model = AutoModelForTokenClassification.from_pretrained(checkpoint)
      # if adapter_args.train_adapter:
          # load_adapter = checkpoint + "/" + args.task_name
          # load_lang_adapter = "{}/{}".format(checkpoint, adapter_args.language)
          # model.model_name = args.model_name_or_path
          # model, lang_adapter_names, task_name = setup_adapter(args, adapter_args, model, load_adapter=load_adapter)
# 
      # model.to(args.device)
      # result, _ = evaluate(args, model, tokenizer, labels, pad_token_label_id, mode="dev", prefix=global_step, lang=args.train_langs, lang2id=lang2id, lang_adapter_names=lang_adapter_names, task_name=task_name, calc_weight_step=args.calc_weight_step)
      # if result["f1"] > best_f1:
        # best_checkpoint = checkpoint
        # best_f1 = result["f1"]
      # if global_step:
        # result = {"{}_{}".format(global_step, k): v for k, v in result.items()}
      # results.update(result)

    output_eval_file = os.path.join(args.output_dir, "eval_results.txt")
    with open(output_eval_file, "w") as writer:
      for key in sorted(results.keys()):
        writer.write("{} = {}\n".format(key, str(results[key])))
      # writer.write("best checkpoint = {}, best f1 = {}\n".format(best_checkpoint, best_f1))
  if args.do_predict and args.local_rank in [-1, 0]:
    model, tokenizer, lang2id = load_model(args, num_labels)
  # Prediction
    logger.info('Evaluating the model on test set of all the languages specified')
    
    #Set up the task adapter
    if adapter_args.train_adapter or args.test_adapter:
        load_adapter = (best_checkpoint + "/" + args.task_name) if args.predict_task_adapter is None else args.predict_task_adapter
        # load_adapter = 'output/panx/bert-base-multilingual-cased-LR1e-4-epoch100-MaxLen128-TrainLangen_en_s0/checkpoint-best/ner/'
        logger.info(f'Task Adapter will be loaded from this path {load_adapter}')
        load_lang_adapter = args.predict_lang_adapter
        model.model_name = args.model_name_or_path
        model, lang_adapter_names, task_name = setup_adapter(args, adapter_args, model, load_adapter=load_adapter, load_lang_adapter=load_lang_adapter)
    model.to(args.device)

    predict_and_save(args, adapter_args, model, tokenizer, labels, lang2id, pad_token_label_id, lang_adapter_names, task_name, 'test')

  if args.do_predict_train and args.local_rank in [-1, 0]:
    logger.info('Evaluating on the train set of all specified languages')
    model, tokenizer, lang2id = load_model(args, num_labels)

    if adapter_args.train_adapter or args.test_adapter:
        load_adapter = (best_checkpoint + "/" + args.task_name) if args.predict_task_adapter is None else args.predict_task_adapter
        # load_adapter = 'output/panx/bert-base-multilingual-cased-LR1e-4-epoch100-MaxLen128-TrainLangen_en_s0/checkpoint-best/ner/'
        logger.info(f'Task Adapter will be loaded from this path {load_adapter}')
        load_lang_adapter = args.predict_lang_adapter
        model.model_name = args.model_name_or_path
        model, lang_adapter_names, task_name = setup_adapter(args, adapter_args, model, load_adapter=load_adapter, load_lang_adapter=load_lang_adapter)
    model.to(args.device)

    predict_and_save(args, adapter_args, model, tokenizer, labels, lang2id, pad_token_label_id, lang_adapter_names, task_name, 'train')
  
  #Predict dev set
  if args.do_predict_dev and args.local_rank in [-1, 0]:
    model, tokenizer, lang2id = load_model(args, num_labels)

    logger.info('Evaluating on the dev sets of all the specified languages')
    
    #Set up task and language adapters
    if adapter_args.train_adapter or args.test_adapter:
        load_adapter = (best_checkpoint + "/" + args.task_name) if args.predict_task_adapter is None else args.predict_task_adapter
        # load_adapter = 'output/panx/bert-base-multilingual-cased-LR1e-4-epoch100-MaxLen128-TrainLangen_en_s0/checkpoint-best/ner/'
        logger.info(f'Task Adapter will be loaded from this path {load_adapter}')
        load_lang_adapter = args.predict_lang_adapter
        model.model_name = args.model_name_or_path
        model, lang_adapter_names, task_name = setup_adapter(args, adapter_args, model, load_adapter=load_adapter, load_lang_adapter=load_lang_adapter)
    model.to(args.device)

    predict_and_save(args, adapter_args, model, tokenizer, labels, lang2id, pad_token_label_id, lang_adapter_names, task_name, 'dev')

def save_predictions(args, predictions, output_file, text_file, idx_file, output_word_prediction=False):
  # Save predictions
  with open(text_file, "r") as text_reader, open(idx_file, "r") as idx_reader:
    text = text_reader.readlines()
    index = idx_reader.readlines()
    assert len(text) == len(index)

  # Sanity check on the predictions
  with open(output_file, "w") as writer:
    example_id = 0
    prev_id = int(index[0])
    for line, idx in zip(text, index):
      if line == "" or line == "\n":
        example_id += 1
      else:
        cur_id = int(idx)
        output_line = '\n' if cur_id != prev_id else ''
        if output_word_prediction:
          output_line += line.split()[0] + '\t'
        output_line += predictions[example_id].pop(0) + '\n'
        writer.write(output_line)
        prev_id = cur_id

if __name__ == "__main__":
  main()
