Add AlignScore.py class of transformer model - easy to use

AlignScoreCS.py

@@ -0,0 +1,634 @@
+import transformers
+from transformers import PretrainedConfig
+import os
+from pathlib import Path
+import numpy as np
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
+import torch.nn as nn
+import torch
+# This include should be add when using different AlignScoreFunction methods instead of score() 
+# from nltk.tokenize import sent_tokenize
+from tqdm import tqdm
+
+class AlignScoreCS(transformers.XLMRobertaModel):
+    """
+    ALIGNSCORE class
+
+    Description:
+        Model ALIGNSCORECS has been trained according the paper for 3 days on 4GPUs AMD NVIDIA.
+        (3 epochs, 1e-5 learning rate, 1e-6 AdamWeps, batchsize 32, WarmupRatio 0.06, 0.1 WeighDecay)
+        - XLMROBERTA-base model with 3 classification HEAD {regression,binary,3way} using shared encoder
+    
+    USAGE: AlignScore.py
+        - from_pretrained - loads the model, usage as transformers.model
+        - .score(context, claim) - function
+                - returns probs of the ALIGNED class using 3way class head as in the paper.
+
+        alignScoreCS = AlignScoreCS.from_pretrained("/mnt/data/factcheck/AlignScore-data/AAmodel/MTLModel/mo
+        alignScoreCS.score(context,claim)
+
+        If you want to try different classification head use parameter:
+            - task_name = "re" : regression head
+            - task_name = "bin" : binary classification head
+            - task_name = "3way" : 3way classification head
+
+    """
+    _regression_model = "re_model"
+    _binary_class_model = "bin_model"
+    _3way_class_model = "3way_model"
+    
+    def __init__(self, encoder, taskmodels_dict, model_name= "xlm-roberta-large", **kwargs):
+        super().__init__(transformers.XLMRobertaConfig(), **kwargs)
+        self.encoder = encoder
+        self.taskmodels_dict = nn.ModuleDict(taskmodels_dict)
+        self.tokenizer = None
+        self.model_name =  model_name
+        self.inferencer = None
+    
+    def init_inferencer(self, device = "cuda"):
+        self.tokenizer = transformers.AutoTokenizer.from_pretrained(self.model_name) if not self.tokenizer else self.tokenizer
+        self.inferencer = self.InferenceHandler(self, self.tokenizer, device)
+    
+    
+    
+    """
+        Score: scores the context and claim with Aligned probabitlity of 3way classification head
+         - using paper code inferencer from ALignScore
+    
+    """
+    def score(self, context, claim, **kwargs):
+        if self.inferencer is None:
+            self.init_inferencer()
+        scores = self.inferencer.nlg_eval(context, claim)
+        return scores
+    
+    """
+        Score: scores the context and claim with ALIGNED probability (wrt task_name ["re" | "bin" | "3way"])
+        
+        Returns the probability of the ALIGNED CLASS between context text and claim text
+         - chunks text by 350 tokens and splits claim into sentences
+         - using 3way classification head
+    """
+    def score_sentences(self, context :str, claim :str, task_name = "3way", batch_size = 2, return_all_outputs = False, **kwargs):
+        self.tokenizer = transformers.AutoTokenizer.from_pretrained(self.model_name) if not self.tokenizer else self.tokenizer
+        chunked_inputs = self.chunk_sent_input(context,claim, chunk_size=350,chunk_claim_size=150)
+        nclaims, ncontexts = (chunked_inputs["n_claims"],chunked_inputs["n_contexts"])
+        with torch.no_grad():
+            chunked_inputs = {key : torch.tensor(item).to(self.device) for key, item in chunked_inputs.items() if not key.startswith("n_")}
+            chunked_outputs = {}
+            for i in range(0,len(chunked_inputs["input_ids"]),batch_size):
+                tmp = self.forward(task_name = task_name,**{"input_ids":chunked_inputs["input_ids"][i:i+batch_size],"attention_mask" :chunked_inputs["attention_mask"][i:i+batch_size]}, **kwargs)
+                for k, item in tmp.items():
+                    chunked_outputs[k] = chunked_outputs.get(k, []) + [item]
+            logits = torch.vstack(chunked_outputs["logits"]).cpu()
+            outputs = {"score" : self.alignscore_input(logits,nclaims=nclaims,ncontexts=ncontexts, task_name=task_name)}
+            outputs["outputs"] = chunked_outputs
+        return torch.tensor([outputs["score"]]) if not return_all_outputs else outputs
+    
+    
+    """
+        Score: scores the context and claim with ALIGNED probability (wrt task_name ["re" | "bin" | "3way"])
+        
+        Returns the probability of the ALIGNED CLASS between context text and claim text
+         - chunks text into 350 tolens and chunks claim into 150 tokens
+         - using 3way classification head
+    """
+    def score_chunks(self, context :str, claim :str, task_name = "3way", batch_size = 2, return_all_outputs = False, **kwargs):
+        self.tokenizer = transformers.AutoTokenizer.from_pretrained(self.model_name) if not self.tokenizer else self.tokenizer
+        chunked_inputs = self.chunk_inputs(context,claim, chunk_size=350)
+        chunked_inputs = {key : torch.tensor(item).to(self.device) for key, item in chunked_inputs.items()}
+        chunked_outputs = self.forward(task_name = task_name, **chunked_inputs, **kwargs)
+        outputs = {"score" : self.alignscore_input_deprecated(chunked_outputs.logits.cpu(), task_name=task_name)}
+        outputs["outputs"] = chunked_outputs
+        return outputs["score"] if not return_all_outputs else outputs
+    
+    """
+        Classify: classify the context and claim to the class label given the task_name ["re" | "bin" | "3way"]
+        
+        Returns the class of {Neutral, contradict, aligned} between context text and claim text
+         - using 3way classification head
+    """
+    def classify(self, context :str, claim :str, task_name = "3way", return_all_outputs = False, **kwargs):
+        self.tokenizer = transformers.AutoTokenizer.from_pretrained(self.model_name) if not self.tokenizer else self.tokenizer
+        chunked_inputs = self.chunk_inputs(context,claim, chunk_size=350)
+        chunked_inputs = {key : torch.tensor(item).to(self.device) for key, item in chunked_inputs.items()}
+        chunked_outputs = self.forward(task_name = task_name, **chunked_inputs, **kwargs)
+        outputs = {"class" : self.get_system_label(chunked_outputs.logits.cpu(), task_name=task_name)}
+        outputs["outputs"] = chunked_outputs
+        return outputs["class"] if not return_all_outputs else outputs
+    
+    
+    def score_truncated(self, context :str, claim :str, task_name = "3way",  return_all_outputs = False, **kwargs):
+        self.tokenizer = transformers.AutoTokenizer.from_pretrained(self.model_name) if not self.tokenizer else self.tokenizer
+        tokenized_inputs = self.tokenizer(list(zip([context], [claim])), padding = "max_length", truncation = True, max_length = 512, return_tensors="pt")
+        tokenized_inputs = {key : torch.tensor(item).to(self.device) for key, item in tokenized_inputs.items()}
+        with torch.no_grad():
+            model_outputs = self.forward(task_name=task_name, **tokenized_inputs, **kwargs)
+            outputs = {"score" : self.alignscore_input(model_outputs["logits"].cpu(),nclaims=1, ncontexts=1, task_name=task_name)}
+            outputs["outputs"] = model_outputs
+        return torch.tensor([outputs["score"]]) if not return_all_outputs else outputs
+    
+    def forward(self, task_name = "3way", **kwargs):
+        return self.taskmodels_dict[task_name](**kwargs)
+    
+    def __call__(self, task_name, **kwargs):
+        return self.taskmodels_dict[task_name](**kwargs)
+    
+    """
+        Get the probability of the ALIGNED label from input
+    """
+    def alignscore_input(self, chunked_logits, nclaims, ncontexts, task_name = "3way"):
+        if task_name == "re":
+            ouptuts = chunked_logits.detach()
+            # Reshape the tensor to separate each block of n rows
+            reshaped_tensor = ouptuts.view(nclaims, ncontexts)
+
+            # Extract the maximum values from the first column (index 0) within each block of n rows
+            max_values, _ = reshaped_tensor.max(dim=1)
+
+            # Calculate the mean of the max values for each block of n rows
+            mean_of_maxes = torch.mean(max_values, dim=0)
+            return mean_of_maxes.tolist()
+        else:
+            nlabels = {"3way" : 3, "re" : 1, "2way" : 2}[task_name]
+            ouptuts = chunked_logits.softmax(1).detach()
+            # Reshape the tensor to separate each block of n rows
+            reshaped_tensor = ouptuts.view(nclaims, ncontexts, nlabels)
+
+            # Extract the maximum values from the first column (index 0) within each block of n rows
+            max_values, _ = torch.max(reshaped_tensor[:, :, 1], dim=1)
+
+            # Calculate the mean of the max values for each block of n rows
+            mean_of_maxes = torch.mean(max_values, dim=0)
+            return mean_of_maxes.tolist()
+        
+        
+    def alignscore_input_deprecated(self, chunked_logits, task_name = "3way"):
+        if task_name == "re":
+            return chunked_logits.detach().amax(0).tolist()
+        else:
+            return chunked_logits.softmax(1).detach()[:, 1].amax(0).tolist() # return max probability over the ALIGNED class
+    
+    
+    """
+        get the label from the input
+    """
+    def get_system_label(self, chunked_logits, task_name):
+        if task_name == "re":
+            return (chunked_logits.sum(0) / chunked_logits.size()[0]).detach().tolist()
+        else:
+            avg_probs = chunked_logits.softmax(1).sum(0) / chunked_logits.size()[0]
+            numpy_array = chunked_logits.softmax(1).argmax(1).detach().numpy()
+            # Calculate the frequencies of each value
+            unique_values, counts = np.unique(numpy_array, return_counts=True)
+            # Find the maximum count
+            max_count = np.max(counts)
+            # Find all values with the maximum count
+            most_frequent_values = unique_values[counts == max_count]
+            return most_frequent_values[0] if most_frequent_values.size == 1 else avg_probs.detach().argmax().tolist()
+
+    """
+        Chunks input context and claim  - context is chunked into 350 tokens
+                                        - claim is chunked into sentences
+                                        - using stride for overflowing tokens 
+    """
+    def chunk_sent_input(self, context, claim, max_length = 512, chunk_size = 350, chunk_claim_size = 150):
+        assert chunk_size <= max_length, "Chunk size {} cannot be greater than max size {}".format(chunk_size, chunk_claim_size, max_length)
+        chunk_claim_size = max_length - chunk_size if chunk_claim_size is None else chunk_claim_size
+        assert chunk_size + chunk_claim_size <= max_length, "Chunk size {} and Chunk claim size {} cannot be together greater than max size {}".format(chunk_size, chunk_claim_size, max_length)
+        return_chunked_inputs = {}
+        context_chunks = self.chunk_text(context, chunk_size=chunk_size, overflowing_tokens_stride = 25, first_special_token=[0])
+        claim_chunks = self.chunk_sentences(claim, chunk_size=chunk_claim_size,overflowing_tokens_stride=int(chunk_claim_size/3), first_special_token=[2])
+        for claim_chunk in claim_chunks:
+            for context_chunk in context_chunks:
+                inputs,attention =self.fill_with_pad_tokens(context_chunk,claim_chunk )
+                return_chunked_inputs["input_ids"] = return_chunked_inputs.get("input_ids",[]) + [inputs]
+                return_chunked_inputs["attention_mask"] = return_chunked_inputs.get("attention_mask",[]) + [attention]
+        return_chunked_inputs["n_claims"] = len(claim_chunks)
+        return_chunked_inputs["n_contexts"] = len(context_chunks)
+        return return_chunked_inputs
+   
+    """
+        According to paper - chunk the text into smaller parts (350tokens + claim_tokens) when the tokenized inputs exceed the max_length
+        returns chunked input
+    """
+    def chunk_inputs(self, context, claim, max_length = 512, chunk_size = 512, first_fit_within_max_length = True):
+        assert chunk_size <= max_length, "Chunk size {} cannot be greater than max size {}".format(chunk_size, max_length)
+        
+        tokenized_claim = self.tokenizer(claim, return_length=True)
+        tokenized_claim["input_ids"][0] = 2 # </s> token according to pair tokenization where the separator of the context and claim is </s></s>
+        tokenized_context = self.tokenizer(context, return_length = True)
+        assert tokenized_claim["length"][0] < max_length*4/5, "Create chunks of claim sentences. Claim is too long {} which is more than 4/5 from {}.".format(tokenized_claim["length"][0], max_length)
+
+        # set chunk size to incorporate the claim size as well
+        chunk_size = min(max_length, chunk_size + tokenized_claim["length"][0])
+        
+        first_check_max_size = max_length if first_fit_within_max_length else chunk_size
+        
+        if tokenized_claim["length"][0] + tokenized_context["length"][0] <= first_check_max_size: #if it fits within max_length
+            input_ids, attention_mask = self.fill_with_pad_tokens(tokenized_context["input_ids"],tokenized_claim["input_ids"])
+            return {"input_ids" : [input_ids], "attention_mask" : [attention_mask]}
+        else: # make chunks
+            return_chunked_inputs = {}
+            current_chunk = {}
+            for sentence in sent_tokenize(context, language="czech"):
+                tok_sent = self.tokenizer(sentence, return_length=True)
+                if len(current_chunk.get("input_ids",[0])) + tok_sent["length"][0] - 1  + tokenized_claim["length"][0] <= chunk_size:
+                    current_chunk["input_ids"] = current_chunk.get("input_ids",[0]) + tok_sent["input_ids"][1:-1]
+                else:
+                    return_chunked_inputs = self._update_chunked_inputs(tokenized_claim, current_chunk, return_chunked_inputs, max_length, tok_sent)
+                    current_chunk["input_ids"] = [0] + tok_sent["input_ids"][1:-1]
+            if current_chunk != {}: # add the rest
+                    return_chunked_inputs = self._update_chunked_inputs(tokenized_claim, current_chunk, return_chunked_inputs, max_length)
+                    current_chunk = {}
+            return return_chunked_inputs
+   
+    """
+        Chunks input context and claim  - context is chunked into 350 tokens
+                                        - claim is chunked into 150 tokens 
+                                        - using stride for overflowing tokens 
+    """
+    def chunk_input_deprecated(self, context, claim, max_length = 512, chunk_size = 350, chunk_claim_size = 150):
+        assert chunk_size <= max_length, "Chunk size {} cannot be greater than max size {}".format(chunk_size, chunk_claim_size, max_length)
+        chunk_claim_size = max_length - chunk_size if chunk_claim_size is None else chunk_claim_size
+        assert chunk_size + chunk_claim_size <= max_length, "Chunk size {} and Chunk claim size {} cannot be together greater than max size {}".format(chunk_size, chunk_claim_size, max_length)
+        return_chunked_inputs = {}
+        context_chunks = self.chunk_text(context, chunk_size=chunk_size, overflowing_tokens_stride = 25, first_special_token=[0])
+        claim_chunks = self.chunk_text(claim, chunk_size=chunk_claim_size,overflowing_tokens_stride=int(chunk_claim_size/3), first_special_token=[2])
+        for claim_chunk in claim_chunks:
+            for context_chunk in context_chunks:
+                inputs,attention =self.fill_with_pad_tokens(context_chunk,claim_chunk )
+                return_chunked_inputs["input_ids"] = return_chunked_inputs.get("input_ids",[]) + [inputs]
+                return_chunked_inputs["attention_mask"] = return_chunked_inputs.get("attention_mask",[]) + [attention]
+        return_chunked_inputs["n_claims"] = len(claim_chunks)
+        return_chunked_inputs["n_contexts"] = len(context_chunks)
+        return return_chunked_inputs
+    
+   
+    """
+        Chunk texts into blocks of chunk_size tokens
+        
+    """ 
+    def chunk_text(self, text, chunk_size = 350, overflowing_tokens_stride = 25, language="czech", first_special_token = [0]):
+        sentences = sent_tokenize(text, language=language)
+        tokenized = self.tokenizer(sentences if sentences != [] else [""], return_length=True)
+        chunks = []
+        chunk, current_chunk_size = ([], 0)
+        for i, length in enumerate(tokenized["length"]):
+            
+            # WRAP THE TOKENIZED SENTNECE INTO LIST TO HANDLE OVERFLOWING TOKENS EASILY
+            # Case when length of one sentence is longer than the chunk size - split the sentence into chunks of chunk size
+            if length > chunk_size:
+                splits = [first_special_token + tokenized["input_ids"][i][max(1,cs):min(cs + chunk_size - 2, length - 1)] + [2] for cs in range(0, length , chunk_size-(2+overflowing_tokens_stride))]
+            # Case when lenght of sequence is equal or smaller than the chunk size - only continue
+            else:
+                splits = [first_special_token + tokenized["input_ids"][i][1:]] 
+            
+            # Go through sentence or splits of sentence
+            for subsentence in splits:
+                up_length = len(subsentence) - 2 
+                
+                # Case when the current chunk = 0
+                if current_chunk_size == 0:
+                    current_chunk_size = up_length + 2  # First include <s> and </s>  tokens
+                    chunk = subsentence[:-1]
+                # Case when the current chunk + length of new subsentence <= chunk_size - only add
+                elif current_chunk_size + up_length <= chunk_size:
+                    current_chunk_size += up_length
+                    chunk += subsentence[1:-1]
+                # Case when the current chunk + length of new subsentence > chunk_size - create chunk
+                else:
+                    chunks += [chunk + [2]]
+                    current_chunk_size = up_length + 2  # First include <s> and </s>  tokens
+                    chunk = subsentence[:-1]
+        #Case when the loop ended but the current chunk isnt saved in the chunks
+        if chunk != []:
+            chunks += [chunk + [2]]
+        # lengths = [len(ch) for ch in chunks]
+        # print("Lenght in tokens of ",len(lengths)," chunks (AVG=",np.mean(lengths),",MAX=",np.max(lengths),",MIN=", np.min(lengths),")")
+        return chunks
+    
+    """
+        Chunks text into sentences using nlt.sent_tokenize
+    """
+    def chunk_sentences(self, text, chunk_size, overflowing_tokens_stride = 0, language="czech", sentence_window = 2, first_special_token = [2]):
+        sentences = sent_tokenize(text, language=language)
+        tokenized = self.tokenizer(sentences if sentences != [] else [""], return_length=True)
+        chunks = []
+        current_chunk = []
+        for i, length in enumerate(tokenized["length"]):
+            # WRAP THE TOKENIZED SENTNECE INTO LIST TO HANDLE OVERFLOWING TOKENS EASILY
+            # Case when length of one sentence is longer than the chunk size - split the sentence into chunks of chunk size
+            if length > chunk_size:
+                splits = [first_special_token + tokenized["input_ids"][i][max(1,cs):min(cs + chunk_size - 2, length - 1)] + [2] for cs in range(0, length , chunk_size-(2+overflowing_tokens_stride))]
+            # Case when lenght of sequence is equal or smaller than the chunk size - only continue
+            else:
+                splits = [first_special_token + tokenized["input_ids"][i][1:]] 
+            
+            #Go through sentence or parts of sentence and create chunks
+            for split in splits:
+                chunks += [split]
+                # if len(current_chunk) == sentence_window:
+                #     chunks += [first_special_token + [item for row in current_chunk for item in row] + [2]]
+                #     current_chunk = current_chunk[1:] + [split[1:-1]]
+                # else:
+                #     current_chunk += [split[1:-1]]
+                    
+        # if chunks == []:
+        #     chunks += [first_special_token + [item for row in current_chunk for item in row] + [2]]
+        return chunks
+
+    """
+        join context and claim tokens as input_ids and create attention_mask    
+    """
+    def fill_with_pad_tokens(self, first, second, max_length=512, pad_token = 1):
+        return first + second + [pad_token]*max(max_length-len(first)-len(second),0), [1]*(len(first)+len(second)) + [0]*max(max_length-len(first)-len(second),0)
+        
+    
+    def _update_chunked_inputs(self, tokenized_claim, current_chunk, return_chunked_inputs, max_length, tok_sent = {"input_ids" : []}):
+        # truncate if there is a long sentence (rare occurrences)
+        if len(current_chunk.get("input_ids",[0])) + tokenized_claim["length"][0] >= max_length:
+            chunk = current_chunk["input_ids"].copy()[:max_length-tokenized_claim["length"][0]-1] + [2]
+        elif not current_chunk.get("input_ids",False):
+            chunk = tok_sent["input_ids"][: max_length - tokenized_claim["length"][0] -1] + [2]
+        else:
+            chunk = current_chunk["input_ids"].copy() + [2] # add </s> end of sentence
+        claim_ids = tokenized_claim["input_ids"].copy()
+        inputs, attention = self.fill_with_pad_tokens(chunk,claim_ids )
+        return_chunked_inputs["input_ids"] = return_chunked_inputs.get("input_ids",[]) + [inputs]
+        return_chunked_inputs["attention_mask"] = return_chunked_inputs.get("attention_mask",[]) + [attention]
+        return return_chunked_inputs
+        
+    @classmethod
+    def get_encoder_attr_name(cls, model):
+        """
+        The encoder transformer is named differently in each model "architecture".
+        This method lets us get the name of the encoder attribute
+        """
+        model_class_name = model.__class__.__name__
+        if model_class_name.startswith("XLMRoberta"):
+            return "roberta"
+        else:
+            raise KeyError(f"Add support for new model {model_class_name}")
+    
+    
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: Optional[Union[str, os.PathLike]],
+        model_name : str = "xlm-roberta-large",
+        *model_args,
+        config: Optional[Union[PretrainedConfig, str, os.PathLike]] = None,
+        cache_dir: Optional[Union[str, os.PathLike]] = None,
+        ignore_mismatched_sizes: bool = False,
+        force_download: bool = False,
+        local_files_only: bool = False,
+        token: Optional[Union[str, bool]] = None,
+        revision: str = "main",
+        use_safetensors: bool = None,
+        **kwargs,
+    ):
+        # Check if the required model directories exist then load it from file
+        if all(os.path.exists(os.path.join(pretrained_model_name_or_path, model_dir)) for model_dir in [cls._3way_class_model, cls._regression_model, cls._binary_class_model]):
+            # assert all(
+                
+            #     for model_dir in [cls._3way_class_model, cls._regression_model, cls._binary_class_model]
+            # ), "Error: Required model directories not found!"
+            
+            # Disable the warning about newly initialized weights
+            transformers.logging.set_verbosity_error()
+
+            shared_encoder = None
+            taskmodels_dict = {}
+            for path_name in [cls._regression_model, cls._binary_class_model, cls._3way_class_model]:
+                task_name = path_name.split("_")[0]
+                
+                # Load the configuration for the task-specific model
+                task_config = transformers.XLMRobertaConfig.from_json_file("{}/{}/config.json".format(pretrained_model_name_or_path,path_name))
+                # Create the task-specific model
+                model = transformers.XLMRobertaForSequenceClassification.from_pretrained(model_name, config=task_config,*model_args,**kwargs)
+                # Load the weights for the task-specific model
+                model.load_state_dict(torch.load("{}/{}/pytorch_model.bin".format(pretrained_model_name_or_path,path_name), map_location=torch.device('cuda' if torch.cuda.is_available() else 'cpu')))
+                # Set the shared encoder to the model's encoder
+                if shared_encoder is None:
+                    shared_encoder = getattr(model, AlignScoreCS.get_encoder_attr_name(model))
+                else:
+                    setattr(model, AlignScoreCS.get_encoder_attr_name(model), shared_encoder)
+                taskmodels_dict[task_name] = model
+
+            # Create the AlignScoreCS with the shared encoder and loaded task-specific models
+            alignScoreCS = AlignScoreCS(encoder=shared_encoder, taskmodels_dict=taskmodels_dict, model_name=model_name)
+         #Try load the model from huggingface hub
+        else:
+            shared_encoder = None
+            taskmodels_dict = {}
+            for model_dir in [cls._regression_model, cls._binary_class_model, cls._3way_class_model]:
+                task_name = model_dir.split("_")[0]
+                config = transformers.XLMRobertaConfig.from_pretrained(f"{pretrained_model_name_or_path}", subfolder=model_dir)
+                model = transformers.XLMRobertaForSequenceClassification.from_pretrained(f"{pretrained_model_name_or_path}",config=config, subfolder=model_dir)
+                if shared_encoder is None:
+                    shared_encoder = getattr(model, AlignScoreCS.get_encoder_attr_name(model))
+                else:
+                    setattr(model, AlignScoreCS.get_encoder_attr_name(model), shared_encoder)
+                taskmodels_dict[task_name] = model
+            alignScoreCS = AlignScoreCS(encoder=shared_encoder, taskmodels_dict=taskmodels_dict, model_name=model_name)
+
+        return alignScoreCS
+    
+    
+    def save_pretrained(
+        self,
+        save_directory: Union[str, os.PathLike],
+        is_main_process: bool = True,
+        state_dict: Optional[dict] = None,
+        save_function: Callable = torch.save,
+        push_to_hub: bool = False,
+        max_shard_size: Union[int, str] = "10GB",
+        safe_serialization: bool = False,
+        variant: Optional[str] = None,
+        token: Optional[Union[str, bool]] = None,
+        save_peft_format: bool = True,
+        **kwargs,
+    ):
+        for task_name, model_type in self.taskmodels_dict.items():
+            model_type.save_pretrained(save_directory = Path(save_directory,task_name+"_model"),
+                                    is_main_process = is_main_process,
+                                    state_dict = state_dict,
+                                    save_function = save_function,
+                                    push_to_hub = push_to_hub,
+                                    max_shard_size = max_shard_size,
+                                    safe_serialization = safe_serialization,
+                                    variant = variant,
+                                    token = token,
+                                    save_peft_format = save_peft_format,
+                                    **kwargs)
+
+# This piece of code is copied from AlignScore github repository
+# if you want to use different nlg_eval_mode you have to fix errors on your own
+    class InferenceHandler:
+        def __init__(self, model, tokenizer, device = "cuda"):
+            self.model = model
+            self.device = device
+            self.tokenizer = tokenizer
+            self.model.to(self.device)
+            self.model.eval()
+            self.batch_size = 32
+            self.nlg_eval_mode = "nli_sp"
+            self.verbose = False
+            self.task_name = "3way"
+            self.softmax = nn.Softmax(dim=-1)
+        
+        def nlg_eval(self, premise, hypo):
+            if isinstance(premise, str) and isinstance(hypo, str):
+                premise = [premise]
+                hypo = [hypo]
+            return self.inference_example_batch(premise, hypo)
+        
+        def inference_example_batch(self, premise: list, hypo: list):
+            """
+            inference a example,
+            premise: list
+            hypo: list
+            using self.inference to batch the process
+
+            SummaC Style aggregation
+            """
+            self.disable_progress_bar_in_inference = True
+            assert len(premise) == len(hypo), "Premise must has the same length with Hypothesis!"
+
+            out_score = []
+            for one_pre, one_hypo in tqdm(zip(premise, hypo), desc="Evaluating", total=len(premise), disable=(not self.verbose)):
+                out_score.append(self.inference_per_example(one_pre, one_hypo))
+            
+            return torch.tensor(out_score)
+        
+        def inference_per_example(self, premise:str, hypo: str):
+            """
+            inference a example,
+            premise: string
+            hypo: string
+            using self.inference to batch the process
+            """
+            def chunks(lst, n):
+                """Yield successive n-sized chunks from lst."""
+                for i in range(0, len(lst), n):
+                    yield ' '.join(lst[i:i + n])
+            
+            premise_sents = sent_tokenize(premise)
+            premise_sents = premise_sents or ['']
+
+            n_chunk = len(premise.strip().split()) // 350 + 1
+            n_chunk = max(len(premise_sents) // n_chunk, 1)
+            premise_sents = [each for each in chunks(premise_sents, n_chunk)]
+
+            hypo_sents = sent_tokenize(hypo)
+
+            premise_sent_mat = []
+            hypo_sents_mat = []
+            for i in range(len(premise_sents)):
+                for j in range(len(hypo_sents)):
+                    premise_sent_mat.append(premise_sents[i])
+                    hypo_sents_mat.append(hypo_sents[j])
+            
+            if self.nlg_eval_mode is not None:
+                if self.nlg_eval_mode == 'nli_sp':
+                    output_score = self.inference(premise_sent_mat, hypo_sents_mat)[:,1] ### use NLI head OR ALIGN head
+                output_score = output_score.view(len(premise_sents), len(hypo_sents)).max(dim=0).values.mean().item() ### sum or mean depends on the task/aspect
+                
+                return output_score
+
+            
+            output_score = self.inference(premise_sent_mat, hypo_sents_mat) ### use NLI head OR ALIGN head
+            output_score = output_score.view(len(premise_sents), len(hypo_sents)).max(dim=0).values.mean().item() ### sum or mean depends on the task/aspect
+
+            return output_score
+
+        def inference(self, premise, hypo, task_name = None):
+            """
+            inference a list of premise and hypo
+
+            Standard aggregation
+            """
+            task_name = self.task_name if task_name is None else task_name
+            if isinstance(premise, str) and isinstance(hypo, str):
+                premise = [premise]
+                hypo = [hypo]
+            
+            batch = self.batch_tokenize(premise, hypo)
+            output_score = []
+
+            for mini_batch in tqdm(batch, desc="Evaluating", disable=not self.verbose or self.disable_progress_bar_in_inference):
+                mini_batch = mini_batch.to(self.device)
+                with torch.no_grad():
+                    model_output = self.model.forward(task_name=task_name, **mini_batch)
+                    model_output = model_output.logits
+                    if task_name == "re":
+                        model_output = model_output.cpu()
+                    else:
+                        model_output = self.softmax(model_output).cpu()
+                output_score.append(model_output[:,:])
+            
+            output_score = torch.cat(output_score)
+            
+            if self.nlg_eval_mode is not None:
+                if self.nlg_eval_mode == 'nli':
+                    output_score_nli = output_score[:,1]
+                    return output_score_nli
+                elif self.nlg_eval_mode == 'bin':
+                    return output_score
+                elif self.nlg_eval_mode == 'reg':
+                    return output_score
+                else:
+                    ValueError("unrecognized nlg eval mode")
+
+            
+            return output_score
+    
+        def batch_tokenize(self, premise, hypo):
+            """
+            input premise and hypos are lists
+            """
+            assert isinstance(premise, list) and isinstance(hypo, list)
+            assert len(premise) == len(hypo), "premise and hypo should be in the same length."
+
+            batch = []
+            for mini_batch_pre, mini_batch_hypo in zip(self.chunks(premise, self.batch_size), self.chunks(hypo, self.batch_size)):
+                try:
+                    mini_batch = self.tokenizer(mini_batch_pre, mini_batch_hypo, truncation='only_first', padding='max_length', max_length=self.tokenizer.model_max_length, return_tensors='pt')
+                except:
+                    print('text_b too long...')
+                    mini_batch = self.tokenizer(mini_batch_pre, mini_batch_hypo, truncation=True, padding='max_length', max_length=self.tokenizer.model_max_length, return_tensors='pt')
+                batch.append(mini_batch)
+
+            return batch
+        
+        def chunks(self, lst, n):
+            """Yield successive n-sized chunks from lst."""
+            for i in range(0, len(lst), n):
+                yield lst[i:i + n]
+        
+        
+
+if __name__ == "__main__":
+    alignScore = AlignScoreCS.from_pretrained("krotima1/AlignScoreCS")
+    alignScore.to("cuda" if torch.cuda.is_available() else "cpu")
+    
+    print("Tomáš miluje Zuzku!", "|",  "Tomáš miluje Petru!",alignScore.score("Tomáš miluje Zuzku!", "Tomáš miluje Petru."))
+    print("Tomáš miluje Zuzku!", "|",  "Tomáš miluje Zuzku!",alignScore.score("Tomáš miluje Zuzku!", "Tomáš miluje Zuzku!"))
+    print("Tomáš miluje Zuzku.", "|",  "Zuzka miluje Tomáše.",alignScore.score("Tomáš miluje Zuzku!", "Zuzka miluje Tomáše."))
+    print("Tomáš miluje Zuzku.", "|",  "Zuzka nemiluje Tomáše.",alignScore.score("Tomáš miluje Zuzku!", "Zuzka nemiluje Tomáše."))
+    print("Tomáš miluje Zuzku.", "|",  "Tomáš nemiluje Zuzku.",alignScore.score("Tomáš miluje Zuzku!", "Tomáš nemiluje Zuzku."))
+    print("Dva chlapi se perou.", "|",  "Je tu bitka.",alignScore.score("Dva chlapi se perou.", "Je tu bitka."))
+    print("Dva chlapi se perou.", "|",  "Je tu láska.",alignScore.score("Dva chlapi se perou.", "Je tu láska."))
+    print("Karel IV. je Otec vlasti. Nechal postavit katedrálu svatého Víta. \n Kdo nechal vystavět katedrálu?", "|", "Byl to Karel.",alignScore.score("Karel IV. je Otec vlasti. Nechal postavit katedrálu svatého Víta.\nKdo nechal vystavět katedrálu?", "Byl to Karel."))
+    print("Karel IV. je Otec vlasti. Nechal postavit katedrálu svatého Víta. \n Kdo nechal vystavět katedrálu?", "|", "Byl to Vít.",alignScore.score("Karel IV. je Otec vlasti. Nechal postavit katedrálu svatého Víta.\nKdo nechal vystavět katedrálu?", "Byl to Vít."))
+    print("Karel IV. je Otec vlasti. Nechal postavit katedrálu svatého Víta. \n Kdo nechal vystavět katedrálu?", "|", "Byla to katedrála.",alignScore.score("Karel IV. je Otec vlasti. Nechal postavit katedrálu svatého Víta.\nKdo nechal vystavět katedrálu?", "Byla to katedrála."))
+    print("Kdo je Karel IV.? Karel IV. je Otec vlasti. Nechal postavit katedrálu svatého Víta.", "|", "Je Otec.",alignScore.score("Kdo je Karel IV.? Karel IV. je Otec vlasti. Nechal postavit katedrálu svatého Víta.", "Je Otec."))
+    print("Kdo je Karel IV.? Karel IV. je Otec vlasti. Nechal postavit katedrálu svatého Víta.", "|", "Je Otec vlasti.",alignScore.score("Kdo je Karel IV.? Karel IV. je Otec vlasti. Nechal postavit katedrálu svatého Víta.", "Je Otec vlasti."))
+    print("Kdo je Karel IV.? Karel IV. je Otec vlasti. Nechal postavit katedrálu svatého Víta.", "|", "Je katedrála svatého Víta.",alignScore.score("Kdo je Karel IV.? Karel IV. je Otec vlasti. Nechal postavit katedrálu svatého Víta.", "Je katedrála svatého Víta."))
+    print("Karkulka šla do lesa. V lese potkala vlka. Vlk ji zkoušel sníst, ale Karkulka se nedala a Vlkovi utekla!", "|", "Karkulka utekla vklovi.",alignScore.score("Karkulka šla do lesa. V lese potkala vlka. Vlk ji zkoušel sníst, ale Karkulka se nedala a Vlkovi utekla!", "Karkulka utekla vklovi."))
+    print("Karkulka šla do lesa. V lese potkala vlka. Vlk ji zkoušel sníst, ale Karkulka se nedala a Vlkovi utekla!", "|", "Karkulka neutekla vklovi.",alignScore.score("Karkulka šla do lesa. V lese potkala vlka. Vlk ji zkoušel sníst, ale Karkulka se nedala a Vlkovi utekla!", "Karkulka neutekla vklovi."))
+    print("Karkulka šla do lesa. V lese potkala vlka. Vlk ji zkoušel sníst, ale Karkulka se nedala a Vlkovi utekla!", "|", "Vlk snědl Karkulku.",alignScore.score("Karkulka šla do lesa. V lese potkala vlka. Vlk ji zkoušel sníst, ale Karkulka se nedala a Vlkovi utekla!", "Vlk snědl karkulku."))
+    print("Karkulka šla do lesa. V lese potkala vlka. Vlk ji zkoušel sníst, ale Karkulka se nedala a Vlkovi utekla!", "|", "Vlk nesnědl Karkulku.",alignScore.score("Karkulka šla do lesa. V lese potkala vlka. Vlk ji zkoušel sníst, ale Karkulka se nedala a Vlkovi utekla!", "Vlk nesnědl karkulku."))
+    print("Karkulka šla do lesa. V lese potkala vlka. Vlk ji zkoušel sníst, ale Karkulka se nedala a Vlkovi utekla!", "|", "Karkulka snědla vlka.",alignScore.score("Karkulka šla do lesa. V lese potkala vlka. Vlk ji zkoušel sníst, ale Karkulka se nedala a Vlkovi utekla!", "Karkulka snědla vlka."))
+    print("Karkulka šla do lesa. V lese potkala vlka. Vlk ji zkoušel sníst, ale Karkulka se nedala a Vlkovi utekla!", "|", "Karkulka dala vlkovi jablko.",alignScore.score("Karkulka šla do lesa. V lese potkala vlka. Vlk ji zkoušel sníst, ale Karkulka se nedala a Vlkovi utekla!", "Karkulka dala vlkovi jablko."))