text: import os
  import sys
  import time
  import pathlib
  import torch
  import matplotlib.pyplot as plt
  from mpl_toolkits.axes_grid1 import make_axes_locatable
  import numpy as np
  import pandas as pd
  
  from collections import defaultdict
  from transformers import AutoTokenizer
  
  from .utils import (
      get_model,
      get_idf_dict,
      bert_cos_score_idf,
      get_bert_embedding,
      lang2model,
      model2layers,
      get_hash,
      cache_scibert,
      sent_encode,
  )
  
  
  __all__ = ["score", "plot_example"]
  
  
  def score(
      cands,
      refs,
      model_type=None,
      num_layers=None,
      verbose=False,
      idf=False,
      device=None,
      batch_size=64,
      nthreads=4,
      all_layers=False,
      lang=None,
      return_hash=False,
      rescale_with_baseline=False,
      baseline_path=None,
  ):
      """
      BERTScore metric.
  
      Args:
          - :param: `cands` (list of str): candidate sentences
          - :param: `refs` (list of str or list of list of str): reference sentences
          - :param: `model_type` (str): bert specification, default using the suggested
                    model for the target langauge; has to specify at least one of
                    `model_type` or `lang`
          - :param: `num_layers` (int): the layer of representation to use.
                    default using the number of layer tuned on WMT16 correlation data
          - :param: `verbose` (bool): turn on intermediate status update
          - :param: `idf` (bool or dict): use idf weighting, can also be a precomputed idf_dict
          - :param: `device` (str): on which the contextual embedding model will be allocated on.
                    If this argument is None, the model lives on cuda:0 if cuda is available.
          - :param: `nthreads` (int): number of threads
          - :param: `batch_size` (int): bert score processing batch size
          - :param: `lang` (str): language of the sentences; has to specify
                    at least one of `model_type` or `lang`. `lang` needs to be
                    specified when `rescale_with_baseline` is True.
          - :param: `return_hash` (bool): return hash code of the setting
          - :param: `rescale_with_baseline` (bool): rescale bertscore with pre-computed baseline
          - :param: `baseline_path` (str): customized baseline file
  
      Return:
          - :param: `(P, R, F)`: each is of shape (N); N = number of input
                    candidate reference pairs. if returning hashcode, the
                    output will be ((P, R, F), hashcode). If a candidate have 
                    multiple references, the returned score of this candidate is 
                    the *best* score among all references.
      """
      assert len(cands) == len(refs), "Different number of candidates and references"
  
      assert lang is not None or model_type is not None, "Either lang or model_type should be specified"
  
      ref_group_boundaries = None
      if not isinstance(refs[0], str):
          ref_group_boundaries = []
          ori_cands, ori_refs = cands, refs
          cands, refs = [], []
          count = 0
          for cand, ref_group in zip(ori_cands, ori_refs):
              cands += [cand] * len(ref_group)
              refs += ref_group
              ref_group_boundaries.append((count, count + len(ref_group)))
              count += len(ref_group)
  
      if rescale_with_baseline:
          assert lang is not None, "Need to specify Language when rescaling with baseline"
  
      if model_type is None:
          lang = lang.lower()
          model_type = lang2model[lang]
      if num_layers is None:
          num_layers = model2layers[model_type]
  
      if model_type.startswith("scibert"):
          tokenizer = AutoTokenizer.from_pretrained(cache_scibert(model_type))
      else:
          tokenizer = AutoTokenizer.from_pretrained(model_type)
  
      model = get_model(model_type, num_layers, all_layers)
      if device is None:
          device = "cuda" if torch.cuda.is_available() else "cpu"
      model.to(device)
  
      if not idf:
          idf_dict = defaultdict(lambda: 1.0)
          # set idf for [SEP] and [CLS] to 0
          idf_dict[tokenizer.sep_token_id] = 0
          idf_dict[tokenizer.cls_token_id] = 0
      elif isinstance(idf, dict):
          if verbose:
              print("using predefined IDF dict...")
          idf_dict = idf
      else:
          if verbose:
              print("preparing IDF dict...")
          start = time.perf_counter()
          idf_dict = get_idf_dict(refs, tokenizer, nthreads=nthreads)
          if verbose:
              print("done in {:.2f} seconds".format(time.perf_counter() - start))
  
      if verbose:
          print("calculating scores...")
      start = time.perf_counter()
      all_preds = bert_cos_score_idf(
          model,
          refs,
          cands,
          tokenizer,
          idf_dict,
          verbose=verbose,
          device=device,
          batch_size=batch_size,
          all_layers=all_layers,
      ).cpu()
  
      if ref_group_boundaries is not None:
          max_preds = []
          for beg, end in ref_group_boundaries:
              max_preds.append(all_preds[beg:end].max(dim=0)[0])
          all_preds = torch.stack(max_preds, dim=0)
  
      use_custom_baseline = baseline_path is not None
      if rescale_with_baseline:
          if baseline_path is None:
              baseline_path = os.path.join(os.path.dirname(__file__), f"rescale_baseline/{lang}/{model_type}.tsv")
          if os.path.isfile(baseline_path):
              if not all_layers:
                  baselines = torch.from_numpy(pd.read_csv(baseline_path).iloc[num_layers].to_numpy())[1:].float()
              else:
                  baselines = torch.from_numpy(pd.read_csv(baseline_path).to_numpy())[:, 1:].unsqueeze(1).float()
  
              all_preds = (all_preds - baselines) / (1 - baselines)
          else:
              print(f"Warning: Baseline not Found for {model_type} on {lang} at {baseline_path}", file=sys.stderr)
  
      out = all_preds[..., 0], all_preds[..., 1], all_preds[..., 2]  # P, R, F
  
      if verbose:
          time_diff = time.perf_counter() - start
          print(f"done in {time_diff:.2f} seconds, {len(refs) / time_diff:.2f} sentences/sec")
  
      if return_hash:
          return tuple(
              [out, get_hash(model_type, num_layers, idf, rescale_with_baseline, use_custom_baseline=use_custom_baseline)]
          )
  
      return out
  
  
  def plot_example(
      candidate,
      reference,
      model_type=None,
      num_layers=None,
      lang=None,
      rescale_with_baseline=False,
      baseline_path=None,
      fname="",
  ):
      """
      BERTScore metric.
  
      Args:
          - :param: `candidate` (str): a candidate sentence
          - :param: `reference` (str): a reference sentence
          - :param: `verbose` (bool): turn on intermediate status update
          - :param: `model_type` (str): bert specification, default using the suggested
                    model for the target langauge; has to specify at least one of
                    `model_type` or `lang`
          - :param: `num_layers` (int): the layer of representation to use
          - :param: `lang` (str): language of the sentences; has to specify
                    at least one of `model_type` or `lang`. `lang` needs to be
                    specified when `rescale_with_baseline` is True.
          - :param: `return_hash` (bool): return hash code of the setting
          - :param: `rescale_with_baseline` (bool): rescale bertscore with pre-computed baseline
          - :param: `fname` (str): path to save the output plot
      """
      assert isinstance(candidate, str)
      assert isinstance(reference, str)
  
      assert lang is not None or model_type is not None, "Either lang or model_type should be specified"
  
      if rescale_with_baseline:
          assert lang is not None, "Need to specify Language when rescaling with baseline"
  
      if model_type is None:
          lang = lang.lower()
          model_type = lang2model[lang]
      if num_layers is None:
          num_layers = model2layers[model_type]
  
      if model_type.startswith("scibert"):
          tokenizer = AutoTokenizer.from_pretrained(cache_scibert(model_type))
      else:
          tokenizer = AutoTokenizer.from_pretrained(model_type)
      model = get_model(model_type, num_layers)
      device = "cuda" if torch.cuda.is_available() else "cpu"
      model.to(device)
  
      idf_dict = defaultdict(lambda: 1.0)
      # set idf for [SEP] and [CLS] to 0
      idf_dict[tokenizer.sep_token_id] = 0
      idf_dict[tokenizer.cls_token_id] = 0
  
      hyp_embedding, masks, padded_idf = get_bert_embedding(
          [candidate], model, tokenizer, idf_dict, device=device, all_layers=False
      )
      ref_embedding, masks, padded_idf = get_bert_embedding(
          [reference], model, tokenizer, idf_dict, device=device, all_layers=False
      )
      ref_embedding.div_(torch.norm(ref_embedding, dim=-1).unsqueeze(-1))
      hyp_embedding.div_(torch.norm(hyp_embedding, dim=-1).unsqueeze(-1))
      sim = torch.bmm(hyp_embedding, ref_embedding.transpose(1, 2))
      sim = sim.squeeze(0).cpu()
  
      # remove [CLS] and [SEP] tokens
      r_tokens = [tokenizer.decode([i]) for i in sent_encode(tokenizer, reference)][1:-1]
      h_tokens = [tokenizer.decode([i]) for i in sent_encode(tokenizer, candidate)][1:-1]
      sim = sim[1:-1, 1:-1]
  
      if rescale_with_baseline:
          if baseline_path is None:
              baseline_path = os.path.join(os.path.dirname(__file__), f"rescale_baseline/{lang}/{model_type}.tsv")
          if os.path.isfile(baseline_path):
              baselines = torch.from_numpy(pd.read_csv(baseline_path).iloc[num_layers].to_numpy())[1:].float()
              sim = (sim - baselines[2].item()) / (1 - baselines[2].item())
          else:
              print(f"Warning: Baseline not Found for {model_type} on {lang} at {baseline_path}", file=sys.stderr)
  
      fig, ax = plt.subplots(figsize=(len(r_tokens), len(h_tokens)))
      im = ax.imshow(sim, cmap="Blues", vmin=0, vmax=1)
  
      # We want to show all ticks...
      ax.set_xticks(np.arange(len(r_tokens)))
      ax.set_yticks(np.arange(len(h_tokens)))
      # ... and label them with the respective list entries
      ax.set_xticklabels(r_tokens, fontsize=10)
      ax.set_yticklabels(h_tokens, fontsize=10)
      ax.grid(False)
      plt.xlabel("Reference (tokenized)", fontsize=14)
      plt.ylabel("Candidate (tokenized)", fontsize=14)
      title = "Similarity Matrix"
      if rescale_with_baseline:
          title += " (after Rescaling)"
      plt.title(title, fontsize=14)
  
      divider = make_axes_locatable(ax)
      cax = divider.append_axes("right", size="2%", pad=0.2)
      fig.colorbar(im, cax=cax)
  
      # Rotate the tick labels and set their alignment.
      plt.setp(ax.get_xticklabels(), rotation=45, ha="right", rotation_mode="anchor")
  
      # Loop over data dimensions and create text annotations.
      for i in range(len(h_tokens)):
          for j in range(len(r_tokens)):
              text = ax.text(
                  j,
                  i,
                  "{:.3f}".format(sim[i, j].item()),
                  ha="center",
                  va="center",
                  color="k" if sim[i, j].item() < 0.5 else "w",
              )
  
      fig.tight_layout()
      if fname != "":
          plt.savefig(fname, dpi=100)
          print("Saved figure to file: ", fname)
      plt.show()
  