# Path: shorttext/utils/transformers.py
# class WrappedBERTEncoder(BERTObject):
#     """ This is the class that encodes sentences with BERT models.
# 
#     For more information, please refer to the following paper:
# 
#     Jacob Devlin, Ming-Wei Chang, Kenton Lee, Kristina Toutanova, "BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding," arXiv:1810.04805 (2018). [`arXiv
#     <https://arxiv.org/abs/1810.04805>`_]
# 
#     """
#     def __init__(
#             self,
#             model=None,
#             tokenizer=None,
#             max_length=48,
#             nbencodinglayers=4,
#             trainable=False,
#             device='cpu'
#     ):
#         """ This is the constructor of the class that encodes sentences with BERT models.
# 
#         :param model: BERT model (default: None, with model `bert-base-uncase` to be used)
#         :param tokenizer: BERT tokenizer (default: None, with model `bert-base-uncase` to be used)
#         :param max_length: maximum number of tokens of each sentence (default: 48)
#         :param nbencodinglayers: number of encoding layers (taking the last layers to encode the sentences, default: 4)
#         :param device: device the language model is stored (default: `cpu`)
#         :type model: str
#         :type tokenizer: str
#         :type max_length: int
#         :type device: str
#         """
#         super(WrappedBERTEncoder, self).__init__(
#             model=model,
#             tokenizer=tokenizer,
#             trainable=trainable,
#             device=device
#         )
#         self.max_length = max_length
#         self.nbencodinglayers = nbencodinglayers
# 
#     def encode_sentences(self, sentences, numpy=False):
#         """ Encode the sentences into numerical vectors, given by a list of strings.
# 
#         It can output either torch tensors or numpy arrays.
# 
#         :param sentences: list of strings to encode
#         :param numpy: output a numpy array if `True`; otherwise, output a torch tensor. (Default: `False`)
#         :return: encoded vectors for the sentences
#         :type sentences: list
#         :type numpy: bool
#         :rtype: numpy.array or torch.Tensor
#         """
#         input_ids = []
#         tokenized_texts = []
# 
#         for sentence in sentences:
#             marked_text = '[CLS]' + sentence + '[SEP]'
# 
#             encoded_dict = self.tokenizer.encode_plus(
#                 sentence,
#                 add_special_tokens=True,
#                 truncation=True,
#                 max_length=self.max_length,
#                 padding='max_length',
#                 return_tensors='pt'
#             )
# 
#             tokenized_texts.append(self.tokenizer.tokenize(marked_text))
#             input_ids.append(encoded_dict['input_ids'])
# 
#         input_ids = torch.cat(input_ids, dim=0)
#         segments_id = torch.LongTensor(np.array(input_ids > 0))
#         input_ids = input_ids.to(self.device)
#         segments_id = segments_id.to(self.device)
# 
#         if self.trainable:
#             output = self.model(input_ids, segments_id)
#             sentences_embeddings = output[1]
#             hidden_state = output[2]
#         else:
#             with torch.no_grad():
#                 output = self.model(input_ids, segments_id)
#                 sentences_embeddings = output[1]
#                 hidden_state = output[2]
# 
#         alllayers_token_embeddings = torch.stack(hidden_state, dim=0)
#         alllayers_token_embeddings = alllayers_token_embeddings.permute(1, 2, 0, 3)  # swap dimensions to [sentence, tokens, hidden layers, features]
#         processed_embeddings = alllayers_token_embeddings[:, :, (self.number_hidden_layers+1-self.nbencodinglayers):, :]
# 
#         token_embeddings = torch.reshape(processed_embeddings, (len(sentences), self.max_length, -1))
# 
#         if numpy:
#             sentences_embeddings = sentences_embeddings.detach().numpy()
#             token_embeddings = token_embeddings.detach().numpy()
# 
#         return sentences_embeddings, token_embeddings, tokenized_texts


from itertools import product
from ...utils.transformers import WrappedBERTEncoder
import numpy as np
import torch
class BERTScorer:
    """ This is the class that compute the BERTScores between sentences. BERTScores
    include recall BERTScores, precision BERTScores, and F1 BERTSscores.
    For more information, please refer to this paper:

    Tianyi Zhang, Varsha Kishore, Felix Wu, Kilian Q. Weinberger, Yoav Artzi,
    "BERTScore: Evaluating Text Generation with BERT," arXiv:1904.09675 (2019). [`arXiv
    <https://arxiv.org/abs/1904.09675>`_]

    """
    def __init__(
            self,
            model=None,
            tokenizer=None,
            max_length=48,
            nbencodinglayers=4,
            device='cpu'
    ):
        """ It is the class that compute the BERTScores between sentences.

        :param model: BERT model (default: None, with model `bert-base-uncase` to be used)
        :param tokenizer: BERT tokenizer (default: None, with model `bert-base-uncase` to be used)
        :param max_length: maximum number of tokens of each sentence (default: 48)
        :param nbencodinglayers: number of encoding layers (taking the last layers to encode the sentences, default: 4)
        :param device: device the language model is stored (default: `cpu`)
        :type model: str
        :type tokenizer: str
        :type max_length: int
        :type device: str
        """

# Path: shorttext/utils/transformers.py
# class WrappedBERTEncoder(BERTObject):
#     """ This is the class that encodes sentences with BERT models.
# 
#     For more information, please refer to the following paper:
# 
#     Jacob Devlin, Ming-Wei Chang, Kenton Lee, Kristina Toutanova, "BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding," arXiv:1810.04805 (2018). [`arXiv
#     <https://arxiv.org/abs/1810.04805>`_]
# 
#     """
#     def __init__(
#             self,
#             model=None,
#             tokenizer=None,
#             max_length=48,
#             nbencodinglayers=4,
#             trainable=False,
#             device='cpu'
#     ):
#         """ This is the constructor of the class that encodes sentences with BERT models.
# 
#         :param model: BERT model (default: None, with model `bert-base-uncase` to be used)
#         :param tokenizer: BERT tokenizer (default: None, with model `bert-base-uncase` to be used)
#         :param max_length: maximum number of tokens of each sentence (default: 48)
#         :param nbencodinglayers: number of encoding layers (taking the last layers to encode the sentences, default: 4)
#         :param device: device the language model is stored (default: `cpu`)
#         :type model: str
#         :type tokenizer: str
#         :type max_length: int
#         :type device: str
#         """
#         super(WrappedBERTEncoder, self).__init__(
#             model=model,
#             tokenizer=tokenizer,
#             trainable=trainable,
#             device=device
#         )
#         self.max_length = max_length
#         self.nbencodinglayers = nbencodinglayers
# 
#     def encode_sentences(self, sentences, numpy=False):
#         """ Encode the sentences into numerical vectors, given by a list of strings.
# 
#         It can output either torch tensors or numpy arrays.
# 
#         :param sentences: list of strings to encode
#         :param numpy: output a numpy array if `True`; otherwise, output a torch tensor. (Default: `False`)
#         :return: encoded vectors for the sentences
#         :type sentences: list
#         :type numpy: bool
#         :rtype: numpy.array or torch.Tensor
#         """
#         input_ids = []
#         tokenized_texts = []
# 
#         for sentence in sentences:
#             marked_text = '[CLS]' + sentence + '[SEP]'
# 
#             encoded_dict = self.tokenizer.encode_plus(
#                 sentence,
#                 add_special_tokens=True,
#                 truncation=True,
#                 max_length=self.max_length,
#                 padding='max_length',
#                 return_tensors='pt'
#             )
# 
#             tokenized_texts.append(self.tokenizer.tokenize(marked_text))
#             input_ids.append(encoded_dict['input_ids'])
# 
#         input_ids = torch.cat(input_ids, dim=0)
#         segments_id = torch.LongTensor(np.array(input_ids > 0))
#         input_ids = input_ids.to(self.device)
#         segments_id = segments_id.to(self.device)
# 
#         if self.trainable:
#             output = self.model(input_ids, segments_id)
#             sentences_embeddings = output[1]
#             hidden_state = output[2]
#         else:
#             with torch.no_grad():
#                 output = self.model(input_ids, segments_id)
#                 sentences_embeddings = output[1]
#                 hidden_state = output[2]
# 
#         alllayers_token_embeddings = torch.stack(hidden_state, dim=0)
#         alllayers_token_embeddings = alllayers_token_embeddings.permute(1, 2, 0, 3)  # swap dimensions to [sentence, tokens, hidden layers, features]
#         processed_embeddings = alllayers_token_embeddings[:, :, (self.number_hidden_layers+1-self.nbencodinglayers):, :]
# 
#         token_embeddings = torch.reshape(processed_embeddings, (len(sentences), self.max_length, -1))
# 
#         if numpy:
#             sentences_embeddings = sentences_embeddings.detach().numpy()
#             token_embeddings = token_embeddings.detach().numpy()
# 
#         return sentences_embeddings, token_embeddings, tokenized_texts


# Path: shorttext/metrics/transformers/bertscore.py
from itertools import product
from ...utils.transformers import WrappedBERTEncoder
import numpy as np
import torch

class BERTScorer:
    """ This is the class that compute the BERTScores between sentences. BERTScores
    include recall BERTScores, precision BERTScores, and F1 BERTSscores.
    For more information, please refer to this paper:

    Tianyi Zhang, Varsha Kishore, Felix Wu, Kilian Q. Weinberger, Yoav Artzi,
    "BERTScore: Evaluating Text Generation with BERT," arXiv:1904.09675 (2019). [`arXiv
    <https://arxiv.org/abs/1904.09675>`_]

    """
    def __init__(
            self,
            model=None,
            tokenizer=None,
            max_length=48,
            nbencodinglayers=4,
            device='cpu'
    ):
        """ It is the class that compute the BERTScores between sentences.

        :param model: BERT model (default: None, with model `bert-base-uncase` to be used)
        :param tokenizer: BERT tokenizer (default: None, with model `bert-base-uncase` to be used)
        :param max_length: maximum number of tokens of each sentence (default: 48)
        :param nbencodinglayers: number of encoding layers (taking the last layers to encode the sentences, default: 4)
        :param device: device the language model is stored (default: `cpu`)
        :type model: str
        :type tokenizer: str
        :type max_length: int
        :type device: str
        """

        self.encoder = WrappedBERTEncoder(