New Gen

.gitattributes

@@ -36,3 +36,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 validators/meter/ufal-robeczech-base_syllable_BPE_validator_1702489033354 filter=lfs diff=lfs merge=lfs -text
 validators/rhyme/distilroberta-base_syllable_BPE_validator_1702665903087 filter=lfs diff=lfs merge=lfs -text
 validators/year/ufal-robeczech-base_BPE_validator_1702393305267 filter=lfs diff=lfs merge=lfs -text
+utils/validators/meter/ufal-robeczech-base_BPE_validator_1704126400265 filter=lfs diff=lfs merge=lfs -text
+utils/validators/rhyme/distilroberta-base_BPE_validator_1704126399565 filter=lfs diff=lfs merge=lfs -text
+utils/validators/year/ufal-robeczech-base_BPE_validator_1702393305267 filter=lfs diff=lfs merge=lfs -text

corpus_capsulated_datasets.py

@@ -0,0 +1,754 @@
+import os
+import json
+import numpy as np
+import torch
+
+from utils.poet_utils import StropheParams, SyllableMaker, TextAnalysis, TextManipulation
+from torch.utils.data import Dataset
+from transformers import PreTrainedTokenizerBase, PreTrainedModel
+#TODO: Maybe replace year of book being written for year Author was born
+class CorpusDatasetPytorch:
+    """Dataset class responsible for data loading.
+    """
+    
+    class RawDataset:
+        """Dataset distributing raw sting data with no preprocessing
+        """
+        def __init__(self, data_file_paths, lower_case:bool = True):
+            """Construct the frame around Raw data generation
+
+            Args:
+                data_file_paths (_type_): list of paths to data files
+                lower_case (bool, optional): if resulting data should be in lowercase. Defaults to True.
+            """
+            self._data_file_paths = data_file_paths
+            self.lower_case = lower_case
+        
+        def gen_files(self):
+            """Get individual opened files
+
+            Yields:
+                _type_: open file object
+            """
+            for filename in self._data_file_paths:
+                 yield open(filename, 'r') 
+                 
+        def get_text(self):
+            """Get lines of text of poetry
+
+            Yields:
+                str: individual verse line
+            """
+            for step,file in enumerate(self.gen_files()):
+                if step % 500 == 0:
+                    print(f"Processing file {step}")
+                datum = json.load(file)
+                for data_line in datum:
+                    for part_line in data_line['body']:
+                        for text_line in part_line:
+                            yield text_line['text'].lower() if self.lower_case else text_line['text']
+                            
+        def get_part(self):
+            """Get strophe of poetry
+
+            Yields:
+                str: 1 strophe of poetry
+            """
+            for step,file in enumerate(self.gen_files()):
+                if step % 500 == 0:
+                    print(f"Processing file {step}")
+                datum = json.load(file)
+                for data_line in datum:
+                    for part_line in data_line['body']:
+                        part = []
+                        for text_line in part_line:
+                            part.append(text_line['text'])
+                        yield "\n".join(part).lower() if self.lower_case else "\n".join(part)
+        
+        def get_body(self):
+            """Get whole poem
+
+            Yields:
+                str: 1 whole poem
+            """
+            for step,file in enumerate(self.gen_files()):
+                if step % 500 == 0:
+                    print(f"Processing file {step}")
+                datum = json.load(file)
+                for data_line in datum:
+                    body = []
+                    for part_line in data_line['body']:
+                        
+                        for text_line in part_line:
+                            body.append(text_line['text'])
+                        body.append("\n")
+                    yield "\n".join(body).lower() if self.lower_case else "\n".join(body)
+    
+    class TextDataset(Dataset):
+        """Dataset of preprocessed verse lines 
+
+        Args:
+            Dataset (_type_): Dataset is child of torch class for better integration with torch and huggingface
+        """
+        
+        def __init__(self, data_file_paths, prompt_length=True, prompt_ending=True, lower_case=True, val_data_rate: float = 0.05, test_data_rate: float = 0.05):
+            """Construct the class our given data files path and store variables
+
+            Args:
+                data_file_paths (_type_):  list of paths to data files
+                prompt_length (bool, optional): If to prompt the syllable count. Defaults to True.
+                prompt_ending (bool, optional): If to prompt verse ending. Defaults to True.
+                lower_case (bool, optional): If the string should be in lowercase. Defaults to True.
+                val_data_rate (float, optional): Amount of data to be left for validation. Defaults to 0.05.
+                test_data_rate (float, optional): Amount of data to be left for validation. Defaults to 0.05.
+            """
+            self._data_file_paths = data_file_paths
+            self.prompt_length = prompt_length
+            self.prompt_ending = prompt_ending
+            self.lower_case = lower_case
+            
+            self.val_data_rate = val_data_rate
+            self.test_data_rate = test_data_rate
+            
+            self.data = []
+            self.validation_data = []
+            self.test_data = []
+         
+         
+        def gen_files(self):
+            """Get individual opened files
+
+            Yields:
+                _type_: open file object
+            """
+            for filename in self._data_file_paths:
+                 yield open(filename, 'r') 
+                 
+        @staticmethod
+        def _vowels_and_endings(raw_text):
+            """Get the verse ending and number of syllables in verse
+
+            Args:
+                raw_text (str): raw verse to analyze
+
+            Returns:
+                tuple: number of syllables, ending syllable
+            """
+            syllabs = SyllableMaker.syllabify(raw_text)
+            vowels = len(syllabs) #INFO: Now counts the number of syllables
+            ending = syllabs[-1]
+            return vowels, ending
+        
+        @staticmethod
+        def _ending_vector(end):
+            """Construct One-hot encoded vector for ending syllable
+
+            Args:
+                end (str): Ending syllable
+
+            Returns:
+                numpy.ndarray: One-hot encoded vector of ending syllable
+            """
+            verse_end_vector = np.zeros(len(StropheParams.ENDS))
+            if end in StropheParams.ENDS[:-1]:
+                verse_end_vector[StropheParams.ENDS.index(end)] = 1
+            else:
+                verse_end_vector[-1] = 1
+            return verse_end_vector
+        
+        @staticmethod
+        def _syllable_line(raw_text):
+            """Construct verse as sequence of syllables
+
+            Args:
+                raw_text (str): raw verse line
+
+            Returns:
+                str: Verse line as sequence of syllables
+            """
+            ending = raw_text[-1] if raw_text[-1] in [',','.','!','?'] else ''
+            return " ".join(SyllableMaker.syllabify(raw_text)) + ending
+        
+        def _construct_line(self, raw_text, metre):
+            """Construct individual content line
+
+            Args:
+                raw_text (str): raw verse line
+
+            Returns:
+                str: Processed verse line with line parameters
+            """
+            syllables = SyllableMaker.syllabify(raw_text)
+            num_str = f"{len(syllables)} # " if self.prompt_length else ""
+            verse_end = f"{syllables[-1]} # " if self.prompt_ending else ""
+            metre_txt = f"{metre} # "
+            return metre_txt + num_str + verse_end  + raw_text
+        
+        def _introduce_phonetics(self, raw_text:str, phonetics):
+            phonetic_text = raw_text
+            for word in phonetics['words']:
+                phonetic_text = phonetic_text.replace(f'{word["token_lc"]}', f'{word["phoebe"]}') if self.lower_case else phonetic_text.replace(f'{word["token"]}', f'{word["phoebe"]}')
+            return phonetic_text
+        
+        def _construct_syllable_line(self, raw_text, metre):
+            """Construct individual content line as sequence of syllables
+
+            Args:
+                raw_text (str): raw verse line
+
+            Returns:
+                str: Processed verse line as sequence of syllables with line parameters
+            """
+            ending = raw_text[-1] if raw_text[-1] in [',','.','!','?'] else ''
+            syllables = SyllableMaker.syllabify(raw_text)
+            num_str = f"{len(syllables)} # " if self.prompt_length else ""
+            verse_end = f"{syllables[-1]} # " if self.prompt_ending else ""
+            metre_txt = f"{metre} # "
+            return  metre_txt+ num_str + verse_end + " ".join(syllables) + ending
+                     
+            
+        def data_text_line_gen(self):
+            """Preprocess and process data for usage
+            """
+            for step,file in enumerate(self.gen_files()):
+                if step % 500 == 0:
+                    print(f"Processing file {step}")
+                datum = json.load(file)
+                for data_line in datum:
+                    for part_line in data_line['body']:
+                        for text_line in part_line:
+                            metre = StropheParams.METER_TRANSLATE.get(text_line["metre"][0]["type"], "N")
+                            
+                            scanned_text = TextManipulation._remove_most_nonchar(text_line['text'], self.lower_case)
+                            
+                            text_line_scanned = self._construct_line(scanned_text, metre)
+                            syllable_line = self._construct_syllable_line(scanned_text, metre)
+                            #phonetic_text = self._introduce_phonetics(scanned_text, text_line)
+                            
+                            num_vowels, verse_end = self._vowels_and_endings(scanned_text)
+                            
+                            # Based on result of random chose proper set. Because data are large enough, will result in wanted split.
+                            rand_split = np.random.rand()
+                            if rand_split > self.val_data_rate + self.test_data_rate: 
+                                self.data.append({
+                                "input_ids" : [text_line_scanned,syllable_line],
+                                "nums": [num_vowels],
+                                "verse_end": verse_end,
+                                "metre": metre
+                                     })
+                            elif rand_split < self.test_data_rate:
+                                self.test_data.append({
+                                "input_ids" : [text_line_scanned,syllable_line],
+                                "nums": [num_vowels],
+                                "verse_end": verse_end,
+                                "metre": metre
+                                     })
+                            else:
+                                self.validation_data.append({
+                                "input_ids" : [text_line_scanned,syllable_line],
+                                "nums": [num_vowels],
+                                "verse_end": verse_end,
+                                "metre": metre
+                                     })
+                            
+            
+        def __len__(self):
+            """Return length of training data
+
+            Returns:
+                int: length of training data
+            """
+            return len(self.data)
+        
+        def __getitem__(self, index):
+            """return indexed item
+
+            Args:
+                index (int): index from where to return
+
+            Returns:
+                dict: dict with indexed data
+            """
+            return self.data[index]
+        
+    class BodyDataset(Dataset):
+        """Dataset of preprocessed strophe
+
+        Args:
+            Dataset (_type_): Dataset is child of torch class for better integration with torch and huggingface
+        """
+        def __init__(self, data_file_paths,
+                     prompt_length=True, prompt_ending=True, prompt_verse=True, verse_len=[4,6], lower_case=True, val_data_rate: float = 0.05, test_data_rate: float = 0.05):
+            """Construct the class our given data files path and store variables
+
+            Args:
+                data_file_paths (_type_): list of paths to data files
+                prompt_length (bool, optional): If to prompt the syllable count. Defaults to True.
+                prompt_ending (bool, optional): If to prompt verse ending. Defaults to True.
+                prompt_verse (bool, optional): If to prompt rhyme schema . Defaults to True.
+                verse_len (list, optional): Considered length of strophe. Defaults to [4,6].
+                lower_case (bool, optional): If the string should be in lowercase. Defaults to True.
+                val_data_rate (float, optional): Amount of data to be left for validation. Defaults to 0.05.
+                test_data_rate (float, optional): Amount of data to be left for validation. Defaults to 0.05.
+            """
+            self._data_file_paths = data_file_paths
+            self.prompt_length = prompt_length
+            self.prompt_ending = prompt_ending
+            self.prompt_verse = prompt_verse
+            self.verse_len = verse_len
+            self.lower_case = lower_case
+            
+            self.val_data_rate = val_data_rate
+            self.test_data_rate = test_data_rate
+            
+            self.data = []
+            self.validation_data = []
+            self.test_data = []
+        
+        def gen_files(self):
+            """Get individual opened files
+
+            Yields:
+                _type_: open file object
+            """
+            for filename in self._data_file_paths:
+                 yield open(filename, 'r')
+                 
+   
+                     
+        
+        def _construct_line(self, raw_text, metre):
+            """Construct individual content line
+
+            Args:
+                raw_text (str): raw verse line
+
+            Returns:
+                str: Processed verse line with line parameters
+            """
+            syllables = SyllableMaker.syllabify(raw_text)
+            num_str = f"{len(syllables)} # " if self.prompt_length else ""
+            verse_end = f"{syllables[-1]} # " if self.prompt_ending else ""
+            metre_txt = f"{metre} # "
+            return  metre_txt + num_str + verse_end  + raw_text
+        
+        def _construct_syllable_line(self, raw_text, metre):
+            """Construct individual content line as sequence of syllables
+
+            Args:
+                raw_text (str): raw verse line
+
+            Returns:
+                str: Processed verse line as sequence of syllables with line parameters
+            """
+            ending = raw_text[-1] if raw_text[-1] in [',','.','!','?'] else ''
+            syllables = SyllableMaker.syllabify(raw_text)
+            num_str = f"{len(syllables)} # " if self.prompt_length else ""
+            verse_end = f"{syllables[-1]} # " if self.prompt_ending else ""
+            metre_txt = f"{metre} # "
+            return metre_txt + num_str + verse_end + " ".join(syllables) + ending
+            
+            
+                                                           
+        def data_body_gen(self):
+            """Preprocess and process data for usage
+            """
+            for step,file in enumerate(self.gen_files()):
+                if step % 500 == 0:
+                    print(f"Processing file {step}")
+                datum = json.load(file)
+                
+                for data_line in datum:
+                    publish_year_text = TextManipulation._year_bucketor(data_line["biblio"]["year"])
+                    publish_year_true = data_line["biblio"]["year"] if TextAnalysis._is_year(data_line["biblio"]["year"]) else 'NaN'
+                    context = ["NO CONTEXT"]
+
+                    for part_line in data_line['body']:                                                        
+                        body = []
+                        body_syllabs = []
+                        rhyme= []
+                        metres = []
+                        i = 0
+                        for text_line in part_line:
+                            
+                            # In rare cases multiple, but from searching only 1 metre per line
+                            metre = StropheParams.METER_TRANSLATE.get(text_line["metre"][0]["type"], "J")
+                            metres +=  [metre]
+                            
+                            rhyme.append(text_line["rhyme"])  
+                            
+                            scanned_text = TextManipulation._remove_most_nonchar(text_line["text"], self.lower_case)
+
+                            body.append(self._construct_line(scanned_text,metre))
+                            body_syllabs.append(self._construct_syllable_line(scanned_text,metre))
+                            
+                            i+=1
+                            
+                            if i in self.verse_len:
+                                
+                                rhyme_str = TextManipulation._rhyme_string(rhyme)
+                                
+                                text = f"# {rhyme_str} # {publish_year_text}\n" + "\n".join(body) + "\n"
+                                syllable_text = f"# {rhyme_str} # {publish_year_text}\n" + "\n".join(body_syllabs) + "\n"
+                                context_text= "\n".join(context)
+                                rand_split = np.random.rand()
+                                if rand_split > self.val_data_rate + self.test_data_rate:
+                                    self.data.append({
+                                    "input_ids" : [text,syllable_text],
+                                    "context_ids" : context_text,
+                                    "year": publish_year_true,
+                                    "rhyme":  rhyme_str,
+                                    "metre_ids" : metres.copy()
+                                     })
+                                elif rand_split < self.test_data_rate:
+                                    self.test_data.append({
+                                    "input_ids" : [text,syllable_text],
+                                    "context_ids" : context_text,
+                                    "year": publish_year_true,
+                                    "rhyme":  rhyme_str,
+                                    "metre_ids" : metres.copy()
+                                     })
+                                else:
+                                    self.validation_data.append({
+                                    "input_ids" : [text,syllable_text],
+                                    "context_ids" : context_text,
+                                    "year": publish_year_true,
+                                    "rhyme":  rhyme_str,
+                                    "metre_ids" : metres.copy()
+                                     })
+                                
+                                if i == max(self.verse_len):
+                                    body = []
+                                    body_syllabs = []
+                                    rhyme = []
+                                    metres = []
+                                    i=0
+                                
+        
+        def __len__(self):
+            """Return length of training data
+
+            Returns:
+                int: length of training data
+            """
+            return len(self.data)
+        
+        def __getitem__(self, index):
+            """return indexed item
+
+            Args:
+                index (int): index from where to return
+
+            Returns:
+                dict: dict with indexed data
+            """
+            return self.data[index]
+        
+    def get_filenames(self):
+        """Get paths of data files
+
+        Returns:
+            list: Paths of data files
+        """
+        data_filenames = os.listdir(self.data_dir)
+        data_by_files = []
+        for filename in data_filenames:
+            file_path = os.path.join(self.data_dir, filename)
+            data_by_files.append(file_path)
+        return data_by_files
+        
+    def load_raw_(self):
+        """Load Raw dataset with raw string data
+        """
+        filenames = self.get_filenames()
+            
+        self.raw_dataset = CorpusDatasetPytorch.RawDataset(filenames, self.lower_case)
+    
+    def load_json_filenames(self, prompt_length, prompt_ending, prompt_verse, verse_len=[4,6], val_data_rate=0.05, test_data_rate=0.05):
+        """Load Verse and Strophe datasets
+
+        Args:
+            prompt_length (bool, optional): If to prompt the syllable count. Defaults to True.
+            prompt_ending (bool, optional): If to prompt verse ending. Defaults to True.
+            prompt_verse (bool, optional): If to prompt rhyme schema . Defaults to True.
+            verse_len (list, optional): Considered length of strophe. Defaults to [4,6].
+            val_data_rate (float, optional): If the string should be in lowercase. Defaults to 0.1.
+        """
+        filenames = self.get_filenames()
+        
+        self.pytorch_dataset_body = CorpusDatasetPytorch.BodyDataset(filenames, prompt_ending=prompt_ending, 
+                                                    prompt_length=prompt_length, prompt_verse=prompt_verse, 
+                                                    verse_len=verse_len, lower_case=self.lower_case, 
+                                                    val_data_rate=val_data_rate, test_data_rate=test_data_rate)
+        self.pytorch_dataset_body.data_body_gen()
+         
+        
+        self.pytorch_dataset_text = CorpusDatasetPytorch.TextDataset(filenames, prompt_ending=prompt_ending, 
+                                                    prompt_length=prompt_length, lower_case=self.lower_case,
+                                                    val_data_rate=val_data_rate, test_data_rate=test_data_rate)
+        
+        self.pytorch_dataset_text.data_text_line_gen()
+        
+        self.val_pytorch_dataset_body = CorpusDatasetPytorch.BodyDataset([])
+        self.val_pytorch_dataset_text = CorpusDatasetPytorch.TextDataset([])
+        
+        self.val_pytorch_dataset_body.data = self.pytorch_dataset_body.validation_data
+        self.val_pytorch_dataset_text.data = self.pytorch_dataset_text.validation_data
+        
+        self.pytorch_dataset_text.validation_data = []
+        self.pytorch_dataset_body.validation_data = []
+        
+        self.test_pytorch_dataset_body = CorpusDatasetPytorch.BodyDataset([])
+        self.test_pytorch_dataset_text = CorpusDatasetPytorch.TextDataset([])
+        
+        self.test_pytorch_dataset_body.data = self.pytorch_dataset_body.test_data
+        self.test_pytorch_dataset_text.data = self.pytorch_dataset_text.test_data
+        
+        self.pytorch_dataset_text.test_data = []
+        self.pytorch_dataset_body.test_data = []
+        
+    def create_empty(self):
+        """Create empty holder for possible load of processed data from file
+        """
+        self.pytorch_dataset_body = CorpusDatasetPytorch.BodyDataset([])
+        self.pytorch_dataset_text = CorpusDatasetPytorch.TextDataset([])
+        self.val_pytorch_dataset_body = CorpusDatasetPytorch.BodyDataset([])
+        self.val_pytorch_dataset_text = CorpusDatasetPytorch.TextDataset([])
+        self.test_pytorch_dataset_body = CorpusDatasetPytorch.BodyDataset([])
+        self.test_pytorch_dataset_text = CorpusDatasetPytorch.TextDataset([])
+        
+        
+    @staticmethod
+    def collate(batch, tokenizer: PreTrainedTokenizerBase ,max_len = 1024, max_context = 1024 ,mask_rate = 0.0, syllables: bool = False, format: str = 'METER_VERSE'):
+        """Process data for usage in LM
+
+        Args:
+            batch (_type_): Batch with selected data points
+            tokenizer (PreTrainedTokenizerBase): tokenizer to tokenize input text
+            max_len (int, optional): Maximum length of tokenization. Defaults to 1024.
+            max_context (int, optional): Maximum length of tokenization of context. Defaults to 1024.
+            mask_rate (float, optional): Rate in with to mask data. Defaults to 0.0.
+            syllables (bool, optional): If to use sequence of syllables as input text. Defaults to False.
+
+        Returns:
+            dict: tokenized and processed to tensors data
+        """
+        index = 1 if syllables else 0
+        
+        tokenizer.model_max_length = max_len
+        if batch[0]['input_ids'][0].startswith("#"):
+            
+            data = [text['input_ids'][index] for text in batch]
+            if format == "BASIC":
+                data =  ["\n".join
+                         (
+                        [line + f" # {datum.splitlines()[1].split()[0]}"
+                         if i==0 else line.split('#')[-1] for i, line in enumerate(datum.splitlines())] 
+                        ) + tokenizer.eos_token  for j, datum in enumerate(data) 
+                         ]
+            elif format == "VERSE_PAR":
+                 data =  ["\n".join
+                         (
+                        [line + f" # {datum.splitlines()[1].split()[0]}"
+                         if i==0 else "#".join(line.split('#')[1:]) for i, line in enumerate(datum.splitlines())] 
+                        ) + tokenizer.eos_token for j, datum in enumerate(data) 
+                         ]
+            else:
+                data = [text['input_ids'][index] + tokenizer.eos_token for text in batch]
+                 
+            tokenized = tokenizer(data,return_tensors='pt', truncation=True, padding=True)
+            input_ids = tokenized['input_ids']
+            attention = tokenized["attention_mask"]
+        
+        else:
+            tokenized = tokenizer([text['input_ids'][index] + tokenizer.eos_token for text in batch],return_tensors='pt', truncation=True, padding=True)
+            input_ids = tokenized['input_ids']
+            attention = tokenized["attention_mask"]
+    
+        
+        nums = None
+        if "nums" in batch[0].keys():
+            nums = torch.tensor(np.asarray([text['nums'] for text in batch], dtype=np.int32), dtype=torch.float32)
+            
+        rhyme=None
+        if "rhyme" in batch[0].keys():
+            rhyme = torch.tensor(np.asarray([TextAnalysis._rhyme_vector(text["rhyme"]) for text in batch], dtype=np.int32), dtype=torch.float32)
+        
+        verse_end = None
+        if "verse_end" in batch[0].keys():       
+            verse_end = torch.tensor(np.asarray([CorpusDatasetPytorch.TextDataset._ending_vector(text["verse_end"]) for text in batch], dtype=np.int32), dtype=torch.float32)
+        
+        year = None
+        if "year" in batch[0].keys():      
+            year = torch.tensor(np.asarray([TextAnalysis._publish_year_vector(text["year"]) for text in batch], dtype=np.int32), dtype=torch.float32)
+            
+        metre = None
+        if "metre" in batch[0].keys():       
+            metre = torch.tensor(np.asarray([TextAnalysis._metre_vector(text["metre"]) for text in batch], dtype=np.int32), dtype=torch.float32)
+        
+        context_ids = None
+        context_attention_mask = None
+        if "context_ids" in batch[0].keys():
+            tokenizer.model_max_length = max_context
+            tokenized_context = tokenizer([text['context_ids'] + tokenizer.eos_token  for text in batch],return_tensors='pt', truncation=True, padding=True)
+            context_ids = tokenized_context['input_ids']
+            context_attention_mask = tokenized_context['attention_mask'] 
+        
+        return {
+            "input_ids": input_ids,
+            "labels": input_ids.type(torch.LongTensor),
+            "attention_mask": attention,
+            "context_ids" : context_ids,
+            "context_attention_mask" : context_attention_mask,
+            "nums" :  nums,
+            "rhyme": rhyme,
+            "verse_end" : verse_end,
+            "year": year,
+            "metre" : metre}
+        
+        
+    @staticmethod
+    def collate_distil(batch, tokenizer: PreTrainedTokenizerBase ,surrogate_model: PreTrainedModel = None,surrogate_model_device=None ,max_len = 1024):
+        tokenizer.model_max_length = max_len
+        tokenized = tokenizer([text['input_ids'][0] + tokenizer.eos_token for text in batch], return_tensors='pt', truncation=True, padding=True)
+        input_ids = tokenized['input_ids']
+        attention = tokenized["attention_mask"]
+        
+        with torch.no_grad():
+            # This is Tuple
+            model_hidden_states = surrogate_model(input_ids=input_ids.to(surrogate_model_device), 
+                                                  attention_mask=attention.to(surrogate_model_device), 
+                                                  labels=input_ids.type(torch.LongTensor).to(surrogate_model_device))['hidden_states']
+        model_hidden_states = [hidden.cpu().detach() for hidden in model_hidden_states]
+        
+        return {
+            "input_ids": input_ids,
+            "labels": input_ids.type(torch.LongTensor),
+            "attention_mask": attention,
+            "to_replicate_states": model_hidden_states
+         }
+        
+    @staticmethod
+    def collate_validator(batch, tokenizer: PreTrainedTokenizerBase,syllables:bool, is_syllable:bool = False,max_len = 512):
+        """Process data for use in LM for metre,rhyme and year prediction
+
+        Args:
+            batch (_type_): Batch with selected data points
+            tokenizer (PreTrainedTokenizerBase): tokenizer to tokenize input text   
+            syllables (bool): If to use sequence of syllables as input text
+            is_syllable (bool, optional): Signal if the preprocessed inputs contain syllable data. Defaults to False.
+            max_len (int, optional): Maximum length of tokenization. Defaults to 1024.
+
+        Returns:
+            dict: tokenized and processed to tensors data
+        """
+        index = 1 if syllables and is_syllable else 0
+        tokenizer.model_max_length = max_len
+        data_ids = ["\n".join(
+            [" ".join(
+                    SyllableMaker.syllabify(line.split('#')[-1])
+                ) + (line[-1] if line[-1] in [',','.','!','?'] else '') if (syllables and not is_syllable and line) else line.split('#')[-1] for line in text['input_ids'][index].splitlines()[1:]] 
+            ) for text in batch ]
+        
+        
+        tokenized = tokenizer(data_ids, return_tensors='pt', truncation=True, padding=True)
+        input_ids = tokenized['input_ids']
+        attention = tokenized["attention_mask"]
+            
+        rhyme=None
+        if "rhyme" in batch[0].keys():
+            rhyme = torch.tensor(np.asarray([TextAnalysis._rhyme_vector(text["rhyme"]) for text in batch], dtype=np.int32), dtype=torch.float32)
+            
+        year_bucket = None
+        year = None
+        if "year" in batch[0].keys():      
+            year_bucket = torch.tensor(np.asarray([TextAnalysis._publish_year_vector(text["year"]) for text in batch], dtype=np.int32), dtype=torch.float32)
+            year = torch.tensor(np.asarray([ [int(text['year'])] if text['year'] != 'NaN' else [0] for text in batch], dtype=np.int32), dtype=torch.float32)
+        
+        return  {
+            "input_ids": input_ids,
+            "attention_mask": attention,
+            "rhyme": rhyme,
+            "metre_ids": None,
+            "year_bucket": year_bucket,
+            'year':year}
+    
+    @staticmethod
+    def collate_meter(batch, tokenizer: PreTrainedTokenizerBase, syllables:bool, is_syllable:bool = False, max_len = 512):
+        index = 1 if syllables and is_syllable else 0
+        tokenizer.model_max_length = max_len
+        data_ids = []
+        metre = []
+        for datum in batch:
+            data_ids += [
+                    " ".join(
+                    SyllableMaker.syllabify(line.split('#')[-1]) 
+                ) + (line[-1] if line[-1] in [',','.','!','?'] else '') if (syllables and not is_syllable and line) else line.split('#')[-1] for line in datum['input_ids'][index].splitlines()[1:]
+                ]
+            if "metre_ids" in batch[0].keys():
+                metre += [TextAnalysis._metre_vector(one_metre) for one_metre in datum['metre_ids']]
+                
+        tokenized = tokenizer(data_ids, return_tensors='pt', truncation=True, padding=True)
+        input_ids = tokenized['input_ids']
+        attention = tokenized["attention_mask"]
+        
+        metre_ids = None
+        if len(metre) > 0:
+            metre_ids = torch.tensor(np.asarray(metre, dtype=np.int32), dtype=torch.float32)
+            
+        return  {
+            "input_ids": input_ids,
+            "attention_mask": attention,
+            "rhyme": None,
+            "metre_ids": metre_ids,
+            "year_bucket": None,
+            "year": None}
+        
+    
+        
+    def __init__(self, data_dir = "PoetGen\corpusCzechVerse-master\ccv", cache_dir='./', 
+                 prompt_length=True, prompt_ending=True, prompt_verse=True, verse_len=[4,6], lower_case=True, val_data_rate=0.05, test_data_rate=0.05):
+        """Construct the Dataloader and create Datasets
+
+        Args:
+            data_dir (str, optional): Path to data. Defaults to "PoetGen\corpusCzechVerse-master\ccv".
+            cache_dir (str, optional): Path where to store processed data. Defaults to './'.
+            prompt_length (bool, optional): If to prompt the syllable count. Defaults to True.
+            prompt_ending (bool, optional): If to prompt verse ending. Defaults to True.
+            prompt_verse (bool, optional): If to prompt rhyme schema. Defaults to True.
+            verse_len (list, optional): Considered length of strophe. Defaults to [4,6].
+            lower_case (bool, optional): If the string should be in lowercase. Defaults to True.
+            val_data_rate (float, optional): Amount of data to be left for validation. Defaults to 0.1.
+        """
+        self.lower_case = lower_case
+        self.data_dir = data_dir
+        if  os.path.isfile(os.path.join(cache_dir, "body_poet_data.json")) and os.path.isfile(os.path.join(cache_dir, "text_poet_data.json")) \
+            and os.path.isfile(os.path.join(cache_dir, "val_body_poet_data.json")) and os.path.isfile(os.path.join(cache_dir, "val_text_poet_data.json")) \
+            and os.path.isfile(os.path.join(cache_dir, "test_body_poet_data.json")) and os.path.isfile(os.path.join(cache_dir, "test_text_poet_data.json")) :
+            self.create_empty()
+            self.pytorch_dataset_body.data =list(json.load( open( os.path.join(cache_dir, "body_poet_data.json"), 'r')))
+            self.pytorch_dataset_text.data =list(json.load( open( os.path.join(cache_dir, "text_poet_data.json"), 'r')))
+            self.val_pytorch_dataset_body.data = list(json.load( open( os.path.join(cache_dir, "val_body_poet_data.json"), 'r')))
+            self.val_pytorch_dataset_text.data = list(json.load( open( os.path.join(cache_dir, "val_text_poet_data.json"), 'r')))
+            self.test_pytorch_dataset_body.data = list(json.load( open( os.path.join(cache_dir, "test_body_poet_data.json"), 'r')))
+            self.test_pytorch_dataset_text.data = list(json.load( open( os.path.join(cache_dir, "test_text_poet_data.json"), 'r')))
+        else:
+            self.load_json_filenames(prompt_length, prompt_ending, prompt_verse, verse_len=verse_len, val_data_rate=val_data_rate, test_data_rate=test_data_rate)
+            json.dump(self.pytorch_dataset_body.data, open( os.path.join(cache_dir, "body_poet_data.json"), 'w+'), indent = 6)
+            json.dump(self.pytorch_dataset_text.data, open( os.path.join(cache_dir, "text_poet_data.json"), 'w+'), indent = 6)
+            json.dump(self.val_pytorch_dataset_body.data, open( os.path.join(cache_dir, "val_body_poet_data.json"), 'w+'), indent = 6)
+            json.dump(self.val_pytorch_dataset_text.data, open( os.path.join(cache_dir, "val_text_poet_data.json"), 'w+'), indent = 6)
+            json.dump(self.test_pytorch_dataset_body.data, open( os.path.join(cache_dir, "test_body_poet_data.json"), 'w+'), indent = 6)
+            json.dump(self.test_pytorch_dataset_text.data, open( os.path.join(cache_dir, "test_text_poet_data.json"), 'w+'), indent = 6)
+            
+        self.load_raw_()
+        
+        
+        
+#if __name__ == "__main__":
+# Line Count
+#    print(len(list(CorpusDatasetPytorch(os.path.abspath(os.path.join(os.path.dirname(__file__), "corpusCzechVerse", "ccv")) ).raw_dataset.get_text())))
+# Strophe Count
+#    print(len(list(CorpusDatasetPytorch(os.path.abspath(os.path.join(os.path.dirname(__file__), "corpusCzechVerse", "ccv")) ).raw_dataset.get_part())))
+# Poem Count
+#    print(len(list(CorpusDatasetPytorch(os.path.abspath(os.path.join(os.path.dirname(__file__), "corpusCzechVerse", "ccv")) ).raw_dataset.get_body())))

simple_generation_player.py

@@ -0,0 +1,195 @@
+import argparse
+import os
+import torch
+import numpy as np
+
+import sys
+
+from transformers import AutoTokenizer, PreTrainedTokenizerBase, PreTrainedTokenizerFast
+from utils.poet_utils import StropheParams, Tokens, TextManipulation, TextAnalysis
+from utils.base_poet_models import PoetModelBase
+from utils.validators import ValidatorInterface
+
+from corpus_capsulated_datasets import CorpusDatasetPytorch
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument("--model_path_full", default='jinymusim/gpt-czech-poet',  type=str, help="Path to Model")
+
+parser.add_argument("--rhyme_model_path_full", default=os.path.abspath(os.path.join(os.path.dirname(__file__), 'utils', 'validators', 'rhyme', 'distilroberta-base_BPE_validator_1704126399565')),  type=str, help="Path to Model")
+parser.add_argument("--metre_model_path_full", default=os.path.abspath(os.path.join(os.path.dirname(__file__), 'utils' ,"validators", 'meter', 'ufal-robeczech-base_BPE_validator_1704126400265')),  type=str, help="Path to Model")
+parser.add_argument("--year_model_path_full", default=os.path.abspath(os.path.join(os.path.dirname(__file__), 'utils' ,"validators", 'year', 'ufal-robeczech-base_BPE_validator_1702393305267')),  type=str, help="Path to Model")
+
+parser.add_argument("--validator_tokenizer_model_rhyme", default='distilroberta-base', type=str, help="Validator tokenizer")
+parser.add_argument("--validator_tokenizer_model_meter", default='ufal/robeczech-base', type=str, help="Validator tokenizer")
+parser.add_argument("--validator_tokenizer_model_year", default='ufal/robeczech-base', type=str, help="Validator tokenizer")
+parser.add_argument("--val_syllables_rhyme", default=False, type=bool, help="Does validator use syllables")
+parser.add_argument("--val_syllables_meter", default=False, type=bool, help="Does validator use syllables")
+parser.add_argument("--val_syllables_year", default=False, type=bool, help="Does validator use syllables")
+
+
+if __name__ == "__main__":
+    args = parser.parse_args([] if "__file__" not in globals() else None)
+    
+_ ,model_rel_name =  os.path.split(args.model_path_full)
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+model = PoetModelBase(args.model_path_full).to(device)
+model.eval()
+
+rhyme_model, meter_model, year_model = None, None, None
+rhyme_model_name, meter_model_name, year_model_name = "", "", ""
+if args.rhyme_model_path_full:
+    rhyme_model: ValidatorInterface = (torch.load(args.rhyme_model_path_full, map_location=torch.device('cpu'))).to(device)
+    rhyme_model.eval()
+    _,  rhyme_model_name = os.path.split(args.rhyme_model_path_full)
+
+if args.metre_model_path_full:
+    meter_model: ValidatorInterface = (torch.load(args.metre_model_path_full, map_location=torch.device('cpu'))).to(device)
+    meter_model.eval()
+    _, meter_model_name = os.path.split(args.metre_model_path_full)
+    
+if args.year_model_path_full:
+    year_model: ValidatorInterface = (torch.load(args.year_model_path_full, map_location=torch.device('cpu'))).to(device)
+    year_model.eval()
+    _,  year_model_name = os.path.split(args.year_model_path_full)
+# Load Rhyme tokenizer
+validator_tokenizer_rhyme: PreTrainedTokenizerBase = None
+if args.validator_tokenizer_model_rhyme:
+    try:
+        validator_tokenizer_rhyme = AutoTokenizer.from_pretrained(args.validator_tokenizer_model_rhyme)
+    except:
+        validator_tokenizer_rhyme: PreTrainedTokenizerBase = PreTrainedTokenizerFast(tokenizer_file=args.validator_tokenizer_model_rhyme)
+        validator_tokenizer_rhyme.eos_token = Tokens.EOS
+        validator_tokenizer_rhyme.eos_token_id = Tokens.EOS_ID
+        validator_tokenizer_rhyme.pad_token = Tokens.PAD
+        validator_tokenizer_rhyme.pad_token_id = Tokens.PAD_ID
+        validator_tokenizer_rhyme.unk_token = Tokens.UNK
+        validator_tokenizer_rhyme.unk_token_id = Tokens.UNK_ID
+        validator_tokenizer_rhyme.cls_token = Tokens.CLS
+        validator_tokenizer_rhyme.cls_token_id = Tokens.CLS_ID
+        validator_tokenizer_rhyme.sep_token = Tokens.SEP
+        validator_tokenizer_rhyme.sep_token_id = Tokens.SEP_ID
+        
+# Load Meter tokenizer
+validator_tokenizer_meter: PreTrainedTokenizerBase = None
+if args.validator_tokenizer_model_meter:
+    try:
+        validator_tokenizer_meter = AutoTokenizer.from_pretrained(args.validator_tokenizer_model_meter, revision='v1.0')
+    except:
+        validator_tokenizer_meter: PreTrainedTokenizerBase = PreTrainedTokenizerFast(tokenizer_file=args.validator_tokenizer_model_meter)
+        validator_tokenizer_meter.eos_token = Tokens.EOS
+        validator_tokenizer_meter.eos_token_id = Tokens.EOS_ID
+        validator_tokenizer_meter.pad_token = Tokens.PAD
+        validator_tokenizer_meter.pad_token_id = Tokens.PAD_ID
+        validator_tokenizer_meter.unk_token = Tokens.UNK
+        validator_tokenizer_meter.unk_token_id = Tokens.UNK_ID
+        validator_tokenizer_meter.cls_token = Tokens.CLS
+        validator_tokenizer_meter.cls_token_id = Tokens.CLS_ID
+        validator_tokenizer_meter.sep_token = Tokens.SEP
+        validator_tokenizer_meter.sep_token_id = Tokens.SEP_ID
+        
+# Load Year tokenizer
+validator_tokenizer_year: PreTrainedTokenizerBase = None
+if args.validator_tokenizer_model_year:
+    try:
+        validator_tokenizer_year = AutoTokenizer.from_pretrained(args.validator_tokenizer_model_year, revision='v1.0')
+    except:
+        validator_tokenizer_year: PreTrainedTokenizerBase = PreTrainedTokenizerFast(tokenizer_file=args.validator_tokenizer_model_year)
+        validator_tokenizer_year.eos_token = Tokens.EOS
+        validator_tokenizer_year.eos_token_id = Tokens.EOS_ID
+        validator_tokenizer_year.pad_token = Tokens.PAD
+        validator_tokenizer_year.pad_token_id = Tokens.PAD_ID
+        validator_tokenizer_year.unk_token = Tokens.UNK
+        validator_tokenizer_year.unk_token_id = Tokens.UNK_ID
+        validator_tokenizer_year.cls_token = Tokens.CLS
+        validator_tokenizer_year.cls_token_id = Tokens.CLS_ID
+        validator_tokenizer_year.sep_token = Tokens.SEP
+        validator_tokenizer_year.sep_token_id = Tokens.SEP_ID
+ 
+# Load LM tokenizers       
+tokenizer: PreTrainedTokenizerBase =  AutoTokenizer.from_pretrained(args.model_path_full)
+
+generation = "BASIC"
+
+def decoder_helper(type, user_input):
+    if type == "BASIC":
+        tokenized = tokenizer.encode(user_input, return_tensors='pt', truncation=True)
+        out = model.model.generate(tokenized.to(device), 
+                                        max_length=512,
+                                        do_sample=True,
+                                        top_k=50,
+                                        eos_token_id = tokenizer.eos_token_id,
+                                        early_stopping=True,
+                                        pad_token_id= tokenizer.pad_token_id)
+        return tokenizer.decode(out.cpu()[0], skip_special_tokens=True)
+    if type=="FORCED":
+        return model.generate_forced(user_input, tokenizer, sample=True, device=device)
+    
+help = f"Current setting is {generation} generating.\nChange it by writing FORCED/BASIC to input. type HELP for HELP.\nType EXIT to exit."
+  
+print("Welcome to simple czech strophe generation.", help)
+  
+while True:
+    
+    user_input = ""
+    while True:
+        curr_line =  input(">").strip()
+        if curr_line == 'EXIT':
+            sys.exit()
+        elif curr_line == "HELP":
+            print(help)
+            continue
+        elif curr_line == "BASIC":
+            print("Changed to BASIC")
+            generation = 'BASIC'
+            continue
+        elif curr_line == "FORCED":
+            print("Changed to FORCED")
+            generation = "FORCED"
+            continue
+        if not curr_line:
+            break
+        user_input +=  curr_line + "\n"
+        
+    user_input = user_input.strip()
+    user_reqs = model.analyze_prompt(user_input)
+    
+    if "RHYME" not in user_reqs.keys() and generation == "BASIC":
+        print("BASIC generation can't work with imputed format.", help)
+        print("User input is substituted for #")
+        user_input = '#'
+    
+    generated_poem:str = decoder_helper(generation, user_input)
+    
+    # Predictions
+    meters = []
+    rhyme_pred = ''
+    year_pred = 0
+    for line in generated_poem.splitlines():
+        # Skip Empty lines
+        if not line.strip(): 
+            break
+        if not (TextManipulation._remove_most_nonchar(line)).strip():
+            break
+        # Validate for Strophe Parameters
+        if TextAnalysis._is_param_line(line):
+            data = CorpusDatasetPytorch.collate_validator([{"input_ids" :[generated_poem]}],tokenizer=validator_tokenizer_rhyme,
+                                                               is_syllable=False, syllables=args.val_syllables_rhyme,
+                                                               max_len=rhyme_model.model.config.max_position_embeddings - 2)
+            rhyme_pred =StropheParams.RHYME[np.argmax(rhyme_model.predict_state(input_ids=data['input_ids'].to(device)).detach().flatten().cpu().numpy())]
+            data = CorpusDatasetPytorch.collate_validator([{"input_ids" :[generated_poem]}],tokenizer=validator_tokenizer_year,
+                                                               is_syllable=False, syllables=args.val_syllables_year,
+                                                               max_len=year_model.model.config.max_position_embeddings - 2)
+            year_pred = round(year_model.predict_state(input_ids=data['input_ids'].to(device)).detach().flatten().cpu().numpy()[0])
+            continue
+        data = CorpusDatasetPytorch.collate_meter([{"input_ids" :["FIRST LINE SKIP!\n" + line]}],tokenizer=validator_tokenizer_meter,
+                                                           is_syllable=False, syllables=args.val_syllables_meter,
+                                                           max_len=meter_model.model.config.max_position_embeddings - 2)
+        meters.append(
+            StropheParams.METER[np.argmax(meter_model.predict_state(input_ids=data['input_ids'].to(device)).detach().flatten().cpu().numpy())]
+        )
+    print(f"REQUESTED: {user_reqs}, GENERATED USING: {generation}\n")
+    print(generated_poem.strip())
+    print(f"PREDICTED: {rhyme_pred}, {year_pred}, {meters}\n\n")

utils/base_poet_models.py

@@ -0,0 +1,689 @@
+from .poet_model_utils import PoetModelInterface
+from .poet_utils import TextAnalysis, StropheParams
+
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers.utils import ModelOutput
+import random
+import torch
+
+class PoetModelFunctionalInterface(PoetModelInterface):
+    """Poet Model Functional Interface. Abstract class with implementation of 
+
+    Args:
+        PoetModelInterface (_type_): Is child of PoetModelInterface for carrying core methods 
+    """
+    def __init__(self, *args, **kwargs) -> None:
+        """ Constructor. As child Class needs to construct Parent
+        """
+        super().__init__(*args, **kwargs)
+        
+    def analyze_prompt(self, prompt) -> dict:
+        """Analysis of users prompt
+
+        Args:
+            prompt (_type_): dict or string, carrying users intent
+
+        Returns:
+            dict: Analysis with users intended input
+        """
+        if isinstance(prompt, dict):
+            return prompt 
+        features_dict = {}  
+        lines = prompt.splitlines()
+        lines = list(map(str.strip, lines))
+        i = 0
+        while i < len(lines):
+            if not lines[i]:
+                lines.pop(i)
+                i-=1
+            i+=1
+        cont_line = 0
+        for line in lines:
+            if TextAnalysis._is_param_line(line):
+                for key, value in TextAnalysis._first_line_analysis(line).items():
+                    features_dict[key] = value
+            else:
+                val = cont_line
+                if "RHYME" in features_dict.keys() and cont_line < len(features_dict['RHYME']):
+                    if features_dict["RHYME"][cont_line] == "A":
+                        val = 0
+                    elif features_dict["RHYME"][cont_line] == "B":
+                        val = 1
+                    elif features_dict["RHYME"][cont_line] == "C":
+                        val = 2
+                    elif features_dict["RHYME"][cont_line] == "D":
+                        val = 3
+                for key, value in TextAnalysis._continuos_line_analysis(line).items():
+                    features_dict[f"{key}_{val}"] = value
+                cont_line += 1
+                
+        return features_dict
+                   
+    def generate_forced(self, prompt, tokenizer: AutoTokenizer, sample: bool = True, format: str = 'METER_VERSE', device= torch.device('cpu'), *args, **kwargs) -> str:
+        """Generate Strophe using the FORCED generation
+
+        Args:
+            prompt (_type_): dict or string of users intended parameters of strophe start
+            tokenizer (AutoTokenizer): tokenizer to be used during generation. Should be model specific.
+            sample (bool, optional): If to sample. Defaults to False.
+            format (str, optional): Format of generation to be used. Should be same as trained on. possible formats: BASIC, VERSE_PAR, METER_VERSE, OLD (DEPRECATED! For old models compatibility only). Defaults to 'METER_VERSE'.
+            device (_type_, optional): Device to generate on. CPU as default. Defaults to torch.device('cpu').
+
+        Returns:
+            str: Generated Strophe
+        """     
+        features_dict_init = self.analyze_prompt(prompt)
+        # If user parameters as dict, list is initialized to carry future verses.
+        if isinstance(prompt, dict):
+            prompt_list = []
+        else:
+            prompt_list = prompt.splitlines()
+        # GENERATE FOR POSSIBLE MISSING POET PARAM
+        token_gen_rhyme = tokenizer.encode("#", return_tensors='pt')
+        if sample:
+            rhyme_line = self.model.generate(token_gen_rhyme.to(device), 
+                                max_new_tokens= 100,
+                                do_sample=True,
+                                top_k=50,
+                                early_stopping=True,
+                                pad_token_id=tokenizer.pad_token_id,
+                                eos_token_id=tokenizer.eos_token_id)
+        else:
+            rhyme_line = self.model.generate(token_gen_rhyme.to(device), 
+                                max_new_tokens= 100,
+                                num_beams=8,
+                                no_repeat_ngram_size=2,
+                                early_stopping=True,
+                                pad_token_id=tokenizer.pad_token_id,
+                                eos_token_id=tokenizer.eos_token_id)
+        rhyme_dec = tokenizer.decode(rhyme_line.cpu()[0], skip_special_tokens=True).splitlines()[0]
+        features_dict= TextAnalysis._first_line_analysis(rhyme_dec)
+        for key, value in features_dict_init.items():
+            features_dict[key] = value
+        # CONSTRUCT BEST INPUT LINE
+        # BACKUP RHYME
+        if "RHYME" not in features_dict.keys():
+            features_dict["RHYME"] = random.choice(StropheParams.RHYME[:-1])
+        #OLD
+        if format == 'OLD':
+            poet_param_str = ""
+            if "RHYME" in features_dict.keys():
+                poet_param_str += features_dict["RHYME"]
+            if "YEAR" in features_dict.keys():
+                poet_param_str += f" # {features_dict['YEAR']}"
+            if 'STROPHE_METER' in features_dict.keys():
+                poet_param_str += f" # {features_dict['STROPHE_METER']}"
+            
+        elif format != 'METER_VERSE':
+            poet_param_str = "# "
+            if "RHYME" in features_dict.keys():
+                poet_param_str += features_dict["RHYME"]
+            if "YEAR" in features_dict.keys():
+                poet_param_str += f" # {features_dict['YEAR']}"
+            if 'STROPHE_METER' in features_dict.keys():
+                poet_param_str += f" # {features_dict['STROPHE_METER']}"
+        # NEW
+        else:
+            poet_param_str = "# "
+            if "RHYME" in features_dict.keys():
+                poet_param_str += features_dict["RHYME"]
+            if "YEAR" in features_dict.keys():
+                poet_param_str += f" # {features_dict['YEAR']}"
+        # REPLACE OR INSERT BASED ON PRESENCE
+        if len(features_dict_init.keys()) == 0: # Wierd Input
+            prompt_list = [poet_param_str]
+        elif len(prompt_list) == 0: # Inputed as Dict
+            prompt_list.append(poet_param_str)
+        elif "RHYME" not in features_dict_init.keys():
+            if "YEAR" in features_dict_init.keys() or 'STROPHE_METER' in features_dict_init.keys(): # Replace the Uncomplete first line 
+                prompt_list[0] = poet_param_str
+            else:
+                prompt_list.insert(0, poet_param_str)
+        else:
+            prompt_list[0] = poet_param_str
+            
+        verse_len = len(features_dict["RHYME"])
+        
+        # Finish possible not completed lines
+        base_prompt_len = len(prompt_list)
+        for i in range(2,base_prompt_len + 1):
+            rhyme_char = 0
+            if features_dict["RHYME"][(i - 2) % len(features_dict["RHYME"])] == "B":
+                rhyme_char = 1
+            elif features_dict["RHYME"][(i - 2) % len(features_dict["RHYME"])] == "C":
+                rhyme_char = 2
+            elif features_dict["RHYME"][(i - 2) % len(features_dict["RHYME"])] == "D":
+                rhyme_char = 3
+            elif features_dict["RHYME"][(i - 2) % len(features_dict["RHYME"])] == "X":
+                rhyme_char = -1
+            
+            token_gen_finish = tokenizer.encode("\n".join(prompt_list[:i]), return_tensors='pt')
+            if sample:
+                finish_line = self.model.generate(token_gen_finish.to(device), 
+                                    max_new_tokens= 100,
+                                    do_sample=True,
+                                    top_k=50,
+                                    early_stopping=True,
+                                    pad_token_id=tokenizer.pad_token_id,
+                                    eos_token_id=tokenizer.eos_token_id)
+            else:
+                finish_line = self.model.generate(token_gen_finish.to(device), 
+                                    max_new_tokens= 100,
+                                    num_beams=8,
+                                    no_repeat_ngram_size=2,
+                                    early_stopping=True,
+                                    pad_token_id=tokenizer.pad_token_id,
+                                    eos_token_id=tokenizer.eos_token_id)
+            decoded = tokenizer.decode(finish_line.cpu()[0], skip_special_tokens=True).splitlines()
+            to_dec = min(i, len(decoded))
+            prompt_list[:to_dec] = decoded[:to_dec] 
+            
+            if to_dec - 1 < len(prompt_list):
+                dec_line = prompt_list[to_dec-1]
+                #OLD
+                if format == 'VERSE_PAR' or format == 'OLD':
+                    if  f"END_{rhyme_char}" not in features_dict.keys() and len(dec_line.split()) > 1 and rhyme_char>=0 and dec_line.count("#") <=1:
+                        features_dict[f'LENGTH_{rhyme_char}'] = dec_line.split()[0]
+                        features_dict[f'END_{rhyme_char}'] = dec_line.split()[1]
+                    elif f"END_{rhyme_char}" not in features_dict.keys() and len(dec_line.split()) > 2 and rhyme_char>=0:
+                        features_dict[f'LENGTH_{rhyme_char}'] = dec_line.split()[0]
+                        features_dict[f'END_{rhyme_char}'] = dec_line.split()[2]            
+                # NEW
+                elif format == 'METER_VERSE':    
+                    if  f"END_{rhyme_char}" not in features_dict.keys() and len(dec_line.split()) > 4 and rhyme_char>=0:
+                        features_dict[f'METER_{rhyme_char}'] = dec_line.split()[0]
+                        features_dict[f'LENGTH_{rhyme_char}'] = dec_line.split()[2]
+                        features_dict[f'END_{rhyme_char}'] = dec_line.split()[4]
+                
+        
+        
+        # Generating 4 verse rhymes
+        has_rep= False
+        has_rep_again = False
+        while len(prompt_list) <= verse_len:
+            j = 0
+            if features_dict["RHYME"][(len(prompt_list) - 1) % len(features_dict["RHYME"])] == "B":
+                j = 1
+            elif features_dict["RHYME"][(len(prompt_list) - 1) % len(features_dict["RHYME"])] == "C":
+                j = 2
+            elif features_dict["RHYME"][(len(prompt_list) - 1) % len(features_dict["RHYME"])] == "D":
+                j = 3
+            elif features_dict["RHYME"][(len(prompt_list) - 1) % len(features_dict["RHYME"])] == "X":
+                j=-1
+            #OLD
+            if format == 'BASIC':
+                line_start = ""
+            elif format == 'OLD':
+                line_start = (f"{features_dict[f'LENGTH_{j}']} " if f"LENGTH_{j}" in features_dict.keys() else "" ) + \
+                        (f" {features_dict[f'END_{j}'] } #" if  f"END_{j}" in features_dict.keys() else "") 
+            elif format == 'VERSE_PAR':
+                line_start = (f"{features_dict[f'LENGTH_{j}']} #" if f"LENGTH_{j}" in features_dict.keys() else "" ) + \
+                        (f" {features_dict[f'END_{j}'] } #" if  f"END_{j}" in features_dict.keys() else "") 
+            else:
+                line_start = (f"{features_dict[f'METER_{j}'] } #" if f"METER_{j}" in features_dict.keys() else "") + \
+                (f" {features_dict[f'LENGTH_{j}']} #" if f"LENGTH_{j}" in features_dict.keys() else "" ) + \
+                (f" {features_dict[f'END_{j}'] } #" if  f"END_{j}" in features_dict.keys() else "") 
+            tokenized_poet_start = tokenizer.encode("\n".join(prompt_list) + "\n" + line_start,  return_tensors='pt')
+            if sample:
+                out_line =  self.model.generate(tokenized_poet_start.to(device), 
+                                max_new_tokens= 100,
+                                do_sample=True,
+                                top_k=50,
+                                early_stopping=True,
+                                pad_token_id=tokenizer.pad_token_id,
+                                eos_token_id=tokenizer.eos_token_id)
+            else:
+                out_line =  self.model.generate(tokenized_poet_start.to(device), 
+                                max_new_tokens= 100,
+                                num_beams=2,
+                                no_repeat_ngram_size=2,
+                                early_stopping=True,
+                                pad_token_id=tokenizer.pad_token_id,
+                                eos_token_id=tokenizer.eos_token_id)
+            decoded_lines = tokenizer.decode(out_line.cpu()[0], skip_special_tokens=True).splitlines()
+            # Repetition catcher
+           
+            # Possible 
+            if len(decoded_lines) <= len(prompt_list) and not(has_rep_again and has_rep):
+                if has_rep:
+                    prompt_list.pop()
+                    has_rep= False
+                    has_rep_again = True
+                else:
+                    has_rep = True
+                continue
+            if has_rep_again and has_rep:
+                decoded_line: str = decoded_lines[-1]
+            else:
+                decoded_line: str = decoded_lines[len(prompt_list)]
+            #OLD
+            if format == 'VERSE_PAR' or format == 'OLD':
+                if  f"END_{j}" not in features_dict.keys() and len(decoded_line.split()) > 1 and j>=0 and decoded_line.count("#") <=1:
+                    features_dict[f'LENGTH_{j}'] = decoded_line.split()[0]
+                    features_dict[f'END_{j}'] = decoded_line.split()[1]
+                elif f"END_{j}" not in features_dict.keys() and len(decoded_line.split()) > 2 and j>=0:
+                    features_dict[f'LENGTH_{j}'] = decoded_line.split()[0]
+                    features_dict[f'END_{j}'] = decoded_line.split()[2]            
+            # NEW
+            elif format == 'METER_VERSE':    
+                if  f"END_{j}" not in features_dict.keys() and len(decoded_line.split()) > 4 and j>=0:
+                    features_dict[f'METER_{j}'] = decoded_line.split()[0]
+                    features_dict[f'LENGTH_{j}'] = decoded_line.split()[2]
+                    features_dict[f'END_{j}'] = decoded_line.split()[4]
+                
+            prompt_list.append(decoded_line)
+        
+        return "\n".join(prompt_list)
+
+    
+class PoetModelBase(PoetModelFunctionalInterface):
+    def __init__(self, pretrainedModel, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+
+        self.model = AutoModelForCausalLM.from_pretrained(pretrainedModel, output_hidden_states=True)
+            
+        model_config = self.model.config
+        self.model_size = 1
+        # Check for Hidden layer size by Attribute Name
+        if hasattr(model_config, "n_embd"):
+            self.model_size = model_config.n_embd
+        elif hasattr(model_config, "hidden_size"):
+            self.model_size = model_config.hidden_size
+        
+        
+    def forward(self, input_ids=None, labels=None, attention_mask=None, *args, **kwargs):
+        outputs = self.model(input_ids=input_ids, labels=labels, attention_mask=attention_mask)
+        
+        return ModelOutput(loss= outputs.loss, model_output=outputs) # {"model_output" : outputs,"loss" : outputs.loss}
+    
+    def save_LM(self, LM_path):
+        self.model.save_pretrained(LM_path, safe_serialization=False)
+        
+        
+class PoetModelAllTasks(PoetModelFunctionalInterface):
+    def __init__(self, pretrainedModel, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+
+        self.model = AutoModelForCausalLM.from_pretrained(pretrainedModel, output_hidden_states=True)
+            
+        model_config = self.model.config
+        self.model_size = 1
+        # Check for Hidden layer size by Attribute Name
+        if hasattr(model_config, "n_embd"):
+            self.model_size = model_config.n_embd
+        elif hasattr(model_config, "hidden_size"):
+            self.model_size = model_config.hidden_size
+            
+        self.vowels_regressor = torch.nn.Linear(self.model_size,1) # Vowel Count
+        self.rhyme_regressor = torch.nn.Linear(self.model_size, len(StropheParams.RHYME)) # Rhyme Type
+        self.verse_endings = torch.nn.Linear(self.model_size, len(StropheParams.ENDS)) # Verse End Syllable
+        self.metre_regressor = torch.nn.Linear(self.model_size,len(StropheParams.METER)) # Meter Type
+        self.year_regressor = torch.nn.Linear(self.model_size,len(StropheParams.YEAR)) # Year Bucket
+        
+        
+    def forward(self, input_ids=None, labels=None, attention_mask=None, nums=None, rhyme=None, verse_end=None, year=None, metre=None, *args, **kwargs):
+        outputs = self.model(input_ids=input_ids, labels=labels, attention_mask=attention_mask)
+        last_hidden = outputs['hidden_states'][-1]
+        vowel_regression = self.vowels_regressor((last_hidden[:,0,:].view(-1, self.model_size)))
+        rhyme_regression = self.rhyme_regressor((last_hidden[:,0,:].view(-1, self.model_size)))
+        verse_end_reg = self.verse_endings((last_hidden[:,0,:].view(-1, self.model_size)))
+        metre_regression = self.metre_regressor((last_hidden[:,0,:].view(-1, self.model_size)))
+        year_regression = self.year_regressor((last_hidden[:,0,:].view(-1, self.model_size)))
+        full_loss = outputs.loss
+        
+        vowel_loss = None
+        if nums is not None:
+            loss_fct = torch.nn.MSELoss()
+            vowel_loss = loss_fct(vowel_regression.view(-1, 1), nums.view(-1, 1))
+            full_loss = full_loss + 0.1*vowel_loss
+            
+        rhyme_loss = None
+        if rhyme is not None:
+            softmaxed = torch.softmax(rhyme_regression, dim=1)
+            loss_fct = torch.nn.CrossEntropyLoss()
+            rhyme_loss = loss_fct(softmaxed, rhyme)
+            full_loss = full_loss + 0.1*rhyme_loss
+            
+        verse_loss = None
+        if verse_end is not None:
+            softmaxed = torch.softmax(verse_end_reg, dim=1)
+            loss_fct = torch.nn.CrossEntropyLoss()
+            verse_loss = loss_fct(softmaxed, verse_end)
+            full_loss = full_loss + 0.1*verse_loss
+            
+        metre_loss = None
+        if metre is not None:
+            softmaxed = torch.softmax(metre_regression, dim=1)
+            loss_fct = torch.nn.CrossEntropyLoss()
+            metre_loss = loss_fct(softmaxed, metre)
+            full_loss = full_loss + 0.1*metre_loss
+        
+        year_loss = None
+        if year is not None:
+            softmaxed = torch.softmax(year_regression, dim=1)
+            loss_fct = torch.nn.CrossEntropyLoss()
+            year_loss = loss_fct(softmaxed, year)
+            full_loss = full_loss + 0.1*year_loss
+            
+        
+        return {"model_output" : outputs,
+                "vowel_regression_output": vowel_regression, 
+                "vowel_regression_loss": vowel_loss,
+                "rhyme_regression_output": rhyme_regression,
+                "rhyme_regression_loss": rhyme_loss,
+                "verse_end_regression_output" : verse_end_reg,
+                "verse_end_regression_loss" : verse_loss,
+                "metre_regression_output" : metre_regression,
+                "metre_regression_loss" : metre_loss,
+                "year_regression_output" : year_regression,
+                "year_regression_loss" : year_loss,
+                "loss": full_loss}
+    
+    def save_LM(self, LM_path):
+        self.model.save_pretrained(LM_path, safe_serialization=False)
+
+from .poet_model_utils import ContextModule
+        
+class PoetModelContextInput(PoetModelFunctionalInterface):
+    def __init__(self, pretrainedModel, context_input_size:int = 2048, block_count:int=3, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        
+        self.model = AutoModelForCausalLM.from_pretrained(pretrainedModel,output_hidden_states=True)
+             
+        model_config = self.model.config
+        self.model_size = -1
+        # Check for Hidden layer size by Attribute Name
+        if hasattr(model_config, "n_embd"):
+            self.model_size = model_config.n_embd
+        elif hasattr(model_config, "hidden_size"):
+            self.model_size = model_config.hidden_size  # Number of Emmbedings taken from config
+        self.context_size = context_input_size
+            
+            
+        self.model.base_model.h.insert(3, ContextModule(block_count, context_input_size, self.model_size, self.model_size))
+        # Because of Inserted Layer, Head Masks don't match => Add 1 more
+        self.model.base_model.config.n_layer += 1 
+        
+        self.rhyme_regressor = torch.nn.Linear(self.model_size, len(StropheParams.RHYME)) # Rhyme Type
+        
+        
+    def forward(self, input_ids=None, labels=None, attention_mask=None, rhyme=None, context_ids=None, context_attention_mask=None,*args, **kwargs):
+        # Inject Context to bypass GPT2Blocks (Can't Forward it)
+        self.model.base_model.h[3].context_ids = context_ids
+        self.model.base_model.h[3].context_attention_mask = context_attention_mask
+        
+        outputs = self.model(input_ids=input_ids, labels=labels, attention_mask=attention_mask)   
+        last_hidden = outputs['hidden_states'][-1]
+        rhyme_regression = self.rhyme_regressor((last_hidden[:,0,:].view(-1, self.model_size)))
+        full_loss = outputs.loss
+             
+        rhyme_loss = None
+        if rhyme is not None:
+            softmaxed = torch.softmax(rhyme_regression, dim=1)
+            loss_fct = torch.nn.CrossEntropyLoss()
+            rhyme_loss = loss_fct(softmaxed, rhyme)
+            full_loss = full_loss + rhyme_loss
+        # Delete the Injection to prevent Dataloss
+        self.model.base_model.h[3].context_ids = None
+        self.model.base_model.h[3].context_attention_mask = None
+        
+        return {"model_output" : outputs,
+                "rhyme_regression_output": rhyme_regression,
+                "rhyme_regression_loss": rhyme_loss,
+                "loss": full_loss}
+    
+    def save_LM(self, LM_path):
+        self.model.save_pretrained(LM_path)
+
+from .poet_model_utils import PoetTypeModule  
+        
+class PoetModelContextYear(PoetModelFunctionalInterface):
+    def __init__(self, pretrainedModel, context_input_size:int = 2048, block_count:int=3, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        
+        self.model = AutoModelForCausalLM.from_pretrained(pretrainedModel, output_hidden_states=True)        
+        
+        model_config = self.model.config
+        self.model_size = -1
+        # Check for Hidden layer size by Attribute Name
+        if hasattr(model_config, "n_embd"):
+            self.model_size = model_config.n_embd
+        elif hasattr(model_config, "hidden_size"):
+            self.model_size = model_config.hidden_size  # Number of Emmbedings taken from config
+        self.context_size = context_input_size
+            
+            
+        self.model.base_model.h.insert(3, ContextModule(block_count, context_input_size, self.model_size, self.model_size))
+        self.model.base_model.h.insert(3, PoetTypeModule(block_count, context_input_size, self.model_size, self.model_size))
+        # Because of Inserted Layer, Head Masks don't match => Add 1 more
+        self.model.base_model.config.n_layer += 2
+        
+        self.rhyme_regressor = torch.nn.Linear(self.model_size, len(StropheParams.RHYME)) # Rhyme Type
+        self.year_regressor = torch.nn.Linear(self.model_size, len(StropheParams.YEAR)) # Year Bucket
+        
+        
+    def forward(self, input_ids=None, labels=None, attention_mask=None, rhyme=None, context_ids=None, context_attention_mask=None, year=None,*args, **kwargs):
+        # Inject Context to bypass GPT2Blocks (Can't Forward it)
+        self.model.base_model.h[3].context_ids = context_ids
+        self.model.base_model.h[3].context_attention_mask = context_attention_mask
+        self.model.base_model.h[3].type_labels = year
+        
+        self.model.base_model.h[4].context_ids = context_ids
+        self.model.base_model.h[4].context_attention_mask = context_attention_mask
+        
+        outputs = self.model(input_ids=input_ids, labels=labels, attention_mask=attention_mask)
+        last_hidden = outputs['hidden_states'][-1]
+        rhyme_regression = self.rhyme_regressor((last_hidden[:,0,:].view(-1, self.model_size)))
+        full_loss = outputs.loss
+             
+        rhyme_loss = None
+        if rhyme is not None:
+            softmaxed = torch.softmax(rhyme_regression, dim=1)
+            loss_fct = torch.nn.CrossEntropyLoss()
+            rhyme_loss = loss_fct(softmaxed, rhyme)
+            full_loss = full_loss + rhyme_loss
+            
+        
+        year_regression = self.year_regressor((last_hidden[:,0,:].view(-1, self.model_size)))    
+        
+        year_loss = None
+        if year is not None:
+            softmaxed = torch.softmax(year_regression, dim=1)
+            loss_fct = torch.nn.CrossEntropyLoss()
+            year_loss = loss_fct(softmaxed, year)
+            full_loss = full_loss + year_loss +  self.model.base_model.h[3].indiv_loss
+        
+        # Delete the Injection to prevent Dataloss
+        self.model.base_model.h[3].context_ids = None
+        self.model.base_model.h[3].context_attention_mask = None
+        self.model.base_model.h[3].type_labels = None
+        # Delete Loss 
+        self.model.base_model.h[3].indiv_loss = None
+        
+        self.model.base_model.h[4].context_ids = None
+        self.model.base_model.h[4].context_attention_mask = None
+        
+        return {"model_output" : outputs,
+                "rhyme_regression_output": rhyme_regression,
+                "rhyme_regression_loss": rhyme_loss,
+                "year_regression_output" : year_regression,
+                "year_loss" : year_loss,
+                "loss": full_loss}
+    
+    def save_LM(self, LM_path):
+        self.model.save_pretrained(LM_path)
+        
+        
+class DistilModel(PoetModelFunctionalInterface):
+    
+    def __init__(self, pretrainedModel, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+
+        self.model = AutoModelForCausalLM.from_pretrained(pretrainedModel, output_hidden_states=True)
+            
+        model_config = self.model.config
+        self.model_size = 1
+        # Check for Hidden layer size by Attribute Name
+        if hasattr(model_config, "n_embd"):
+            self.model_size = model_config.n_embd
+        elif hasattr(model_config, "hidden_size"):
+            self.model_size = model_config.hidden_size
+        
+        self.kept_states = [1, 3, 5, 7, 9, 11]
+            
+        for pop_index in sorted(list(set(range(len(self.model.base_model.h))) - set(self.kept_states)), reverse=True):
+            
+            self.model.base_model.h.pop(pop_index)
+        # Because of Inserted Layer, Head Masks don't match => Add 1 more
+        self.model.base_model.config.n_layer = len(self.kept_states)
+        
+        self.loss_fnc = torch.nn.MSELoss()
+        
+    def forward(self, input_ids=None, labels=None, attention_mask=None, to_replicate_states= None, *args, **kwargs):
+        outputs = self.model(input_ids=input_ids, labels=labels, attention_mask=attention_mask)
+        loss = outputs.loss
+        # The 6 layers + embeddings (add + 1 to shift the original_index)
+        for distil_index, original_index in enumerate([-1] + self.kept_states):
+            loss += self.loss_fnc(outputs['hidden_states'][distil_index], to_replicate_states[original_index + 1])
+        
+        return {"model_output" : outputs,
+                "loss": loss}
+    
+    def save_LM(self, LM_path):
+        self.model.save_pretrained(LM_path, safe_serialization=False)
+        
+    def generate_forced(self, *args, **kwargs):
+        raise NotImplementedError("Currently without")
+    
+class PoetModelHalfBase(PoetModelFunctionalInterface):
+    def __init__(self, pretrainedModel, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+    
+        self.model = AutoModelForCausalLM.from_pretrained(pretrainedModel, output_hidden_states=True, torch_dtype=torch.float16)
+            
+        model_config = self.model.config
+        self.model_size = -1
+        # Check for Hidden layer size by Attribute Name
+        if hasattr(model_config, "n_embd"):
+            self.model_size = model_config.n_embd
+        elif hasattr(model_config, "hidden_size"):
+            self.model_size = model_config.hidden_size
+        
+        
+    def forward(self, input_ids=None, labels=None, attention_mask=None, *args, **kwargs):
+        outputs = self.model(input_ids=input_ids, labels=labels, attention_mask=attention_mask)
+        
+        return {"model_output" : outputs,
+                "loss" : outputs.loss}
+    
+    def save_LM(self, LM_path):
+        self.model.save_pretrained(LM_path)
+        
+        
+class PoetModelSecondaryTasks(PoetModelFunctionalInterface):
+    def __init__(self, pretrainedModel, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        
+        self.model = AutoModelForCausalLM.from_pretrained(pretrainedModel, output_hidden_states=True)
+            
+        model_config = self.model.config
+        self.model_size = -1
+        # Check for Hidden layer size by Attribute Name
+        if hasattr(model_config, "n_embd"):
+            self.model_size = model_config.n_embd
+        elif hasattr(model_config, "hidden_size"):
+            self.model_size = model_config.hidden_size  # Number of Emmbedings taken from config
+        self.vowels_regressor = torch.nn.Linear(self.model_size,1) # Vowel count
+        self.rhyme_regressor = torch.nn.Linear(self.model_size, len(StropheParams.RHYME)) # Rhyme Type
+        
+        
+    def forward(self, input_ids=None, labels=None, attention_mask=None, nums=None, rhyme=None, *args, **kwargs):
+        outputs = self.model(input_ids=input_ids, labels=labels, attention_mask=attention_mask)
+        last_hidden = outputs['hidden_states'][-1]
+        vowel_regression = self.vowels_regressor((last_hidden[:,0,:].view(-1, self.model_size))) 
+        rhyme_regression = self.rhyme_regressor((last_hidden[:,0,:].view(-1, self.model_size)))
+        full_loss = outputs.loss
+        
+        vowel_loss = None
+        if nums is not None:
+            loss_fct = torch.nn.MSELoss()
+            vowel_loss = loss_fct(vowel_regression.view(-1, 1), nums.view(-1, 1))
+            full_loss = full_loss + vowel_loss
+            
+        rhyme_loss = None
+        if rhyme is not None:
+            softmaxed = torch.softmax(rhyme_regression, dim=1)
+            loss_fct = torch.nn.CrossEntropyLoss()
+            rhyme_loss = loss_fct(softmaxed, rhyme)
+            full_loss = full_loss + rhyme_loss
+            
+        
+        return {"model_output" : outputs,
+                "vowel_regression_output": vowel_regression, 
+                "vowel_regression_loss": vowel_loss,
+                "rhyme_regression_output": rhyme_regression,
+                "rhyme_regression_loss": rhyme_loss,
+                "loss": full_loss}
+    
+    def save_LM(self, LM_path):
+        self.model.save_pretrained(LM_path)
+        
+        
+class PoetModelVerseEnd(PoetModelFunctionalInterface):
+    def __init__(self, pretrainedModel, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        
+        self.model = AutoModelForCausalLM.from_pretrained(pretrainedModel, output_hidden_states=True)
+            
+        model_config = self.model.config
+        self.model_size = -1
+        # Check for Hidden layer size by Attribute Name
+        if hasattr(model_config, "n_embd"):
+            self.model_size = model_config.n_embd
+        elif hasattr(model_config, "hidden_size"):
+            self.model_size = model_config.hidden_size  # Number of Emmbedings taken from config
+        self.vowels_regressor = torch.nn.Linear(self.model_size,1) # Vowel count
+        self.rhyme_regressor = torch.nn.Linear(self.model_size, len(StropheParams.RHYME)) # Rhyme Type
+        self.verse_endings = torch.nn.Linear(self.model_size, len(StropheParams.ENDS)) # Verse End Syllable
+        
+        
+    def forward(self, input_ids=None, labels=None, attention_mask=None, nums=None, rhyme=None, verse_end = None, *args, **kwargs):
+        outputs = self.model(input_ids=input_ids, labels=labels, attention_mask=attention_mask)
+        last_hidden = outputs['hidden_states'][-1]
+        vowel_regression = self.vowels_regressor((last_hidden[:,0,:].view(-1, self.model_size)))
+        rhyme_regression = self.rhyme_regressor((last_hidden[:,0,:].view(-1, self.model_size)))
+        verse_end_reg = self.verse_endings((last_hidden[:,0,:].view(-1, self.model_size)))
+        full_loss = outputs.loss
+        
+        vowel_loss = None
+        if nums is not None:
+            loss_fct = torch.nn.MSELoss()
+            vowel_loss = loss_fct(vowel_regression.view(-1, 1), nums.view(-1, 1))
+            full_loss = full_loss + vowel_loss
+            
+        rhyme_loss = None
+        if rhyme is not None:
+            softmaxed = torch.softmax(rhyme_regression, dim=1)
+            loss_fct = torch.nn.CrossEntropyLoss()
+            rhyme_loss = loss_fct(softmaxed, rhyme)
+            full_loss = full_loss + rhyme_loss
+            
+        verse_loss = None
+        if verse_end is not None:
+            softmaxed = torch.softmax(verse_end_reg, dim=1)
+            loss_fct = torch.nn.CrossEntropyLoss()
+            verse_loss = loss_fct(softmaxed, verse_end)
+            full_loss = full_loss + verse_loss
+            
+        
+        return {"model_output" : outputs,
+                "vowel_regression_output": vowel_regression, 
+                "vowel_regression_loss": vowel_loss,
+                "rhyme_regression_output": rhyme_regression,
+                "rhyme_regression_loss": rhyme_loss,
+                "verse_end_regression_output" : verse_end_reg,
+                "verse_end_regression_loss" : verse_loss,
+                "loss": full_loss}
+    
+    def save_LM(self, LM_path):
+        self.model.save_pretrained(LM_path)

utils/poet_model_utils.py

@@ -0,0 +1,272 @@
+import torch
+
+class PoetModelInterface(torch.nn.Module):
+    """Pytorch Model Interface. Abstract class for all Poet model types
+
+    Args:
+        torch (_type_): Is child of torch.nn.Module for integration with torch and huggingface 
+    """
+    def __init__(self, *args, **kwargs) -> None:
+        """ Constructor. As child Class needs to construct Parent
+        """
+        super().__init__(*args, **kwargs)
+        
+        
+    def forward(self, input_ids=None, labels=None, attention_mask=None, *args, **kwargs):
+        """Compute model output and model loss
+
+        Args:
+            input_ids (_type_, optional): Model inputs. Defaults to None.
+            labels (_type_, optional): Language Model labels. Defaults to None.
+            attention_mask (_type_, optional): Attention mask where padding starts. Defaults to None.
+
+        Raises:
+            NotImplementedError: Abstract class
+        """
+        raise NotImplementedError()
+    
+    def generate_forced(self,  *args, **kwargs):
+        """Generates model output with restriction on inputs and past generation
+
+        Raises:
+            NotImplementedError: Abstract class
+        """
+        raise NotImplementedError()
+
+    @staticmethod
+    def rhyme_like(rhyme:str):
+        """DEPRECATED: Check string in rhyme format 
+
+        Args:
+            rhyme (str): String with possible rhyme
+
+        Returns:
+            bool: Boolean if string like rhyme
+        """
+        return rhyme.isupper() and len(rhyme) in [4,6]
+    
+    def save_LM(self, LM_path):
+        """Save raw LM
+
+        Args:
+            LM_path (str): Where to store the LM
+
+        Raises:
+            NotImplementedError: Abstract class
+        """
+        raise NotImplementedError()
+    
+
+from transformers import GPT2Config, GPT2Model
+from .poet_utils import StropheParams
+
+class ContextModule(torch.nn.Module):
+    """Module for understanding poet context
+
+    Args:
+        torch (_type_): Is child of torch.nn.Module for integration with torch and huggingface 
+    """
+    def __init__(self, block_count, input_size, n_embd ,output_size,*args, **kwargs) -> None:
+        """Construct the underlying small LM for context
+
+        Args:
+            block_count (_type_): LM number of blocks of GPT2Block
+            input_size (_type_): LM size of input
+            n_embd (_type_): LM size of hidden layers
+            output_size (_type_): LM size of output
+        """
+        super().__init__(*args, **kwargs)
+        self.config = GPT2Config(n_positions=input_size, n_head=(n_embd//(768//12)),n_embd=n_embd, 
+                                 n_layer=block_count, output_hidden_states=True,  output_attentions =True)
+        self.context_model = GPT2Model(self.config)
+        self.linear_downscale = torch.nn.Linear(n_embd, output_size)
+        self.input_size = input_size
+        self.n_embd = n_embd
+        self.output_size = output_size
+        # Context is getting injected from Outside
+        self.context_ids = None
+        self.context_attention_mask = None
+    
+    
+    def forward(self, hidden_states,layer_past=None,*args, **kwargs):
+        """Compute Context LM output, Data are injected from outside
+
+        Args:
+            hidden_states (_type_): Current hidden states
+            layer_past (_type_, optional): Past layer outputs. Defaults to None.
+
+        Returns:
+            _type_: GPT2Block structured output (hidden states, layer past, attention, keys)
+        """
+        down = torch.zeros_like(hidden_states)
+        model_output = None
+        # Sometimes there might be no context
+        if self.context_ids != None:
+            model_output = self.context_model.forward(input_ids=self.context_ids, attention_mask=self.context_attention_mask)
+            # Take only the Class token as 
+            down = self.linear_downscale.forward(model_output["hidden_states"][-1][:,0,:].view(-1, self.n_embd))[:, None, :]
+        return  (hidden_states + down,
+                 down[None, :, :, :],
+                 (None if model_output == None else model_output["attentions"], 
+                None))
+        
+class PoetTypeModule(torch.nn.Module):
+    """Module to classify poet type
+
+    Args:
+        torch (_type_): Is child of torch.nn.Module for integration with torch and huggingface 
+    """
+    
+    def __init__(self, block_count, input_size, n_embd,output_size,*args, **kwargs) -> None:
+        """Construct LM for poet classification from inputs
+
+        Args:
+            block_count (_type_): LM number of blocks of GPT2Block
+            input_size (_type_): LM size of input
+            n_embd (_type_): LM size of hidden layers
+            output_size (_type_): LM size of output
+        """
+        super().__init__(*args, **kwargs)
+        self.config = GPT2Config(n_positions=input_size, n_head=(n_embd//(768//12)),n_embd=n_embd, 
+                                 n_layer=block_count, output_hidden_states=True,  output_attentions =True)
+        self.type_model = GPT2Model(self.config)
+        self.type_predict = torch.nn.Linear(n_embd, len(StropheParams.YEAR))
+        self.softmax = torch.nn.Softmax()
+        self.linear_scale = torch.nn.Linear(len(StropheParams.YEAR), output_size)
+        self.input_size = input_size
+        self.n_embd = n_embd
+        self.output_size = output_size
+        # Context and labels are getting injected from Outside
+        self.context_ids = None
+        self.context_attention_mask = None
+        self.type_labels=None
+        # Store for loss for model itself
+        self.indiv_loss=None
+    
+    def forward(self, hidden_states,layer_past=None,*args, **kwargs):
+        """Compute Classification LM output and loss 
+
+        Args:
+            hidden_states (_type_): Current hidden states
+            layer_past (_type_, optional): Past layer outputs. Defaults to None.
+
+        Returns:
+            _type_: GPT2Block structured output (hidden states, layer past, attention, keys)
+        """
+        type_prob = torch.zeros((hidden_states.shape[0], len(StropheParams.YEAR))).to("cuda" if torch.cuda.is_available() else "cpu")
+        model_output = None
+        # Sometimes there might be no context
+        if self.context_ids != None:
+            model_output = self.type_model.forward(input_ids=self.context_ids, attention_mask=self.context_attention_mask)
+            # Only Class token is taken
+            poet_type = self.type_predict.forward(model_output["hidden_states"][-1][:,0,:].view(-1, self.n_embd))
+            type_prob = self.softmax.forward(poet_type) 
+        # If type labels are present, inject the true labels to future blocks
+        if self.type_labels != None:
+            loss_fct = torch.nn.CrossEntropyLoss()
+            self.indiv_loss = loss_fct(type_prob, self.type_labels)
+            type_prob = (self.type_labels.type(torch.FloatTensor)).to("cuda" if torch.cuda.is_available() else "cpu")
+        linear_up = self.linear_scale.forward(type_prob)
+        return (hidden_states + linear_up[:, None, :],
+                linear_up[None, :, None, :], 
+                (None if model_output == None else model_output["attentions"], 
+                None))
+            
+from transformers import PreTrainedTokenizerBase
+
+class ModelManipulation:
+    """Static Class incorporating methods for Manipulation with LMs
+    Code Inspired by article: Fine-tuning the English GPT-2 in any language with Hugging Face
+    Link: https://github.com/piegu/fastai-projects/blob/master/finetuning-English-GPT2-any-language-Portuguese-HuggingFace-fastaiv2.ipynb
+    """
+     
+    @staticmethod
+    def exchange_embedding(poet_model: PoetModelInterface, new_tokenizer: PreTrainedTokenizerBase, old_tokenizer: PreTrainedTokenizerBase, mirror_imbed:bool=False):
+        """Exchange embedding matrixes for GPT2 Models
+
+        Args:
+            poet_model (PoetModelInterface): Model to manipulate with
+            new_tokenizer (PreTrainedTokenizerBase): New tokenization
+            old_tokenizer (PreTrainedTokenizerBase): Old tokenization
+        """
+        # Get old Embeddings
+        if hasattr(poet_model.model, "transformer"):
+            old_embed_in = poet_model.model.transformer.get_input_embeddings().weight.clone().detach()
+        else:
+            old_embed_in = poet_model.model.get_input_embeddings().weight.clone().detach()
+        old_mean_in = old_embed_in.mean(0)
+        # Generate new Embedding based on new tokenization
+        new_embd_in = old_embed_in.new_zeros(new_tokenizer.vocab_size, old_embed_in.size(1))
+        old_vocab = old_tokenizer.get_vocab()
+        
+        vocab_hit = 0
+        # Keep as much from old Embeddings as possible
+        for w, idx_new in new_tokenizer.get_vocab().items():
+            idx_old = old_vocab.get(w, -1)
+            if idx_old >= 0:
+                new_embd_in[idx_new] = old_embed_in[idx_old]
+                vocab_hit +=1
+            else:
+                new_embd_in[idx_new] = old_mean_in
+                
+        print(f"Vocab hit rate: {vocab_hit}/{old_tokenizer.vocab_size}")
+        #Exchange Embeddings and Decoding
+        new_embd_layer_in = torch.nn.Embedding(new_tokenizer.vocab_size, old_embed_in.size(1))
+        new_embd_layer_in.weight.data = new_embd_in
+        if hasattr(poet_model.model, "transformer"):
+            poet_model.model.transformer.set_input_embeddings(new_embd_layer_in)
+        else:
+            poet_model.model.set_input_embeddings(new_embd_layer_in)
+        
+        new_decoder = torch.nn.Linear( old_embed_in.size(1), new_tokenizer.vocab_size, bias=False)
+        if hasattr(poet_model.model, "transformer"):
+            new_decoder.weight = poet_model.model.transformer.wte.weight
+        else:
+            new_decoder.weight = poet_model.model.base_model.embeddings.weight
+        if hasattr(poet_model.model, "lm_head"):
+            poet_model.model.lm_head = new_decoder
+        else:
+            poet_model.model.head = new_decoder
+
+        
+        # Update LM config to reflect possible change in vocab size
+        poet_model.model.config.vocab_size = new_tokenizer.vocab_size
+        
+        
+    @staticmethod
+    def exchange_embedding_roberta(metre_model, new_tokenizer: PreTrainedTokenizerBase, old_tokenizer: PreTrainedTokenizerBase):
+        """Exchange embedding matrixes for Roberta Models
+
+        Args:
+            poet_model (PoetModelInterface): Model to manipulate with
+            new_tokenizer (PreTrainedTokenizerBase): New tokenization
+            old_tokenizer (PreTrainedTokenizerBase): Old tokenization
+        """
+        # Get old Embeddings
+        old_embed = metre_model.model.get_input_embeddings().weight.clone().detach()
+        old_mean = old_embed.mean(0)
+        # Generate new Embedding based on new tokenization
+        new_embd = old_embed.new_zeros(new_tokenizer.vocab_size, old_embed.size(1))
+        old_vocab = old_tokenizer.get_vocab()
+        
+        vocab_hit = 0
+        # Keep as much from old Embeddings as possible
+        for w, idx_new in new_tokenizer.get_vocab().items():
+            idx_old = old_vocab.get(w, -1)
+            if idx_old >= 0:
+                new_embd[idx_new] = old_embed[idx_old]
+                vocab_hit +=1
+            else:
+                new_embd[idx_new] = old_mean
+                
+        print(f"Vocab hit rate: {vocab_hit}/{old_tokenizer.vocab_size}")
+        #Exchange Embeddings and Decoding
+        new_embd_layer = torch.nn.Embedding(new_tokenizer.vocab_size, old_embed.size(1))
+        new_embd_layer.weight.data = new_embd
+        metre_model.model.set_input_embeddings(new_embd_layer)
+        new_decoder = torch.nn.Linear( old_embed.size(1), new_tokenizer.vocab_size)
+        new_decoder.weight = metre_model.model.roberta.embeddings.word_embeddings.weight
+        metre_model.model.lm_head.decoder = new_decoder
+        # Update LM config to reflect possible change in vocab size
+        metre_model.model.config.vocab_size = new_tokenizer.vocab_size
+        

utils/poet_utils.py

@@ -0,0 +1,591 @@
+class StropheParams:
+    
+
+    # Most Common Rhyme Schemas (Every Rhyme schema with presence over 0.36 %)
+    RHYME_SCHEMES = ['ABAB', 'XXXX',
+                 'XAXA','AABB', 
+                 'XXXXXX','ABBA', 
+                 'AAXX', 'AABBCC',
+                 'ABABCC','ABABXX',
+                 'AABCCB','XXAA', 
+                 'XAAX', 'AXAX', 
+                 'XAXAXX','XXABAB',
+                 'ABBACC','AXAA', 
+                 'XAABBX','AABCBC',
+                 'AABBXX','ABBAXX',
+                 'ABABAB','AAXA', 
+                 'AXXA','XAXABB',
+                 'XXAABB','XXAAXX',
+                 'ABABAX','XXABBA',
+                 'AAXBBX','XXXAXA',
+                 'AAAX','XABABX',
+                 'XABBAX','AAXXBB',
+                 'AXABBX','ABABBX',
+                 'XAAXBB','AAAA',
+                 'XAAA','XAABXB',
+                 'AXABXB','AXAXBB',
+                  None]
+    
+    RHYME = RHYME_SCHEMES
+
+
+    NORMAL_SCHEMES = ["ABAB", "ABBA", "AABB", "AABBCC", "ABABCC", "ABBACC", "ABBAAB"]
+
+    # First 200 Most common endings
+    VERSE_ENDS = ['ní', 'la', 'je', 'tí', 'ce', 'ti', 'ky', 'ku', 'li', 'jí', 'ně', 'né', 'vá', 'se', 'ny', 'ly', 'na', 'ne', 'nou', 
+              'lo', 'ci', 'mi', 'ný', 'sti', 'ka', 'le', 'cí', 'ná', 'ží', 'čí', 'ho', 'dí', 'ší', 'du', 'lí', 'dy', 'nu', 'ří', 
+              'ji', 'ru', 'tě', 'ře', 'stí', 'vy', 'ká', 'še', 'dá', 'ni', 'te', 'ví', 'mu', 'tu', 'ta', 'vé', 'val', 'va', 'lý', 
+              'tá', 'že', 'ty', 'no', 'vu', 'lá', 'kem', 'chu', 'ků', 'bě', 'vý', 'sy', 'me', 'zí', 'hu', 'vě', 'lu', 'da', 'ry', 
+              'rá', 'lé', 'ko', 'ři', 'de', 'hy', 'lem', 'tem', 'kou', 'vou', 'ši', 'há', 'sí', 'ze', 'be', 'ra', 'má', 'to', 'by', 
+              'mě', 'su', 'té', 'si', 'ných', 'den', 'či', 'ký', 'ním', 'če', 'tý', 'ma', 'my', 'sem', 'nem', 'dě', 'ha', 'vat', 'ným', 
+              'dem', 'dou', 'sta', 'dla', 'svět', 'zem', 'jen', 'dal', 'mí', 'hou', 'zas', 'sen', 'rem', 'nů', 'bu', 'e', 'ba', 'ké', 
+              'til', 'jest', 'ství', 'děl', 'květ', 'tů', 'chem', 'lou', 'sám', 'bí', 'tou', 'dé', 'šel', 'nul', 'chá', 'vem', 'sa', 
+              'hlas', 'pí', 'čas', 'dil', 'let', 'cích', 'lů', 'žil', 'mů', 'dál', 'cha', 'byl', 'nost', 'ček', 'zy', 'hý', 'nám', 'di', 
+              'bou', 'tím', 'ži', 'tek', 'vil', 'jsem', 'sů', 'dech', 'men', 'tla', 'sá', 'zrak', 'chy', 'vám', 'vi', 'dý', 'rád', 'svou', 
+              'ném', 've', 'py', 'vo', 'vým', 'nek', 'již', 'víc', 'kal', 'mé', 'dů', 'stá', 'dnes', 'sty', 'ven', None]
+    ENDS = VERSE_ENDS
+    # Years to bucket to
+    POET_YEARS_BUCKETS = [1800, 1820, 1840, 1860, 1880, 1900, 1920, 1940, 1960, None]
+    POET_YEARS = POET_YEARS_BUCKETS
+    YEAR = POET_YEARS_BUCKETS
+    # Possible Meter Types
+    METER_TYPES = ["J","T","D","A","X","Y","N","H","P", None]
+    METER = METER_TYPES
+    # Translation of Meter to one char types
+    METER_TRANSLATE = {
+        "J":"J",
+        "T":"T",
+        "D":"D",
+        "A":"A",
+        "X":"X",
+        "Y":"Y",
+        "hexameter": "H",
+        "pentameter": "P",
+        "N":"N"
+    }
+    
+    # Basic Characters to consider in rhyme and syllables (43)
+    VALID_CHARS = [""," ",'a','á','b','c','č','d','ď','e','é','ě',
+               'f','g','h','i','í','j','k','l','m','n','ň',
+               'o','ó','p','q','r','ř','s','š','t','ť','u',
+               'ú','ů','v','w','x','y','ý','z','ž']
+    CHARS = VALID_CHARS
+class Tokens:
+# Tokenizers Special Tokens
+    EOS = "<|EOS|>"
+    EOS_ID = 0
+    PAD = "<|PAD|>"
+    PAD_ID = 1
+    UNK = "<|UNK|>"
+    UNK_ID = 2
+    CLS = "<|CLS|>"
+    CLS_ID = 3
+    # SEP Token is EOS Token
+    SEP = EOS
+    SEP_ID = 0
+    
+    ALL_TOKENS = {
+        EOS : 0,
+        PAD : 1,
+        UNK : 2,
+        CLS : 3,
+    }
+
+
+
+import re
+import numpy as np
+
+def parse_boolean(value):
+    value = value.lower()
+
+    if value in ["true", "yes", "y", "1", "t"]:
+        return True
+    elif value in ["false", "no", "n", "0", "f"]:
+        return False
+
+    return False
+
+class TextManipulation:
+    """Static class for string manipulation methods
+
+    Returns:
+        _type_: str returned by all methods
+    """
+        
+    @staticmethod
+    def _remove_most_nonchar(raw_text, lower_case=True):
+        """Remove most non-alpha non-whitespace characters
+
+        Args:
+            raw_text (str): Text to manipulate
+            lower_case (bool, optional): If resulting text should be lowercase. Defaults to True.
+
+        Returns:
+            str: Cleaned up text
+        """
+        text = re.sub(r'[–\„\“\’\;\:()\]\[\_\*\‘\”\'\-\—\"]+', "", raw_text)
+        return text.lower() if lower_case else text
+    
+    @staticmethod
+    def _remove_all_nonchar(raw_text):
+        """Remove all possible non-alpha characters
+
+        Args:
+            raw_text (str): Text to manipulate
+
+        Returns:
+            str: Cleaned up text
+        """
+        sub = re.sub(r'([^\w\s]+|[0-9]+)', '', raw_text)
+        return sub
+    
+    @staticmethod
+    def _year_bucketor(raw_year):
+        """Bucketizes year string to boundaries, Bad inputs returns NaN string
+
+        Args:
+            raw_year (str): Year string to bucketize
+
+        Returns:
+            _type_: Bucketized year string
+        """
+        if TextAnalysis._is_year(raw_year) and raw_year != "NaN":
+            year_index = np.argmin(np.abs(np.asarray(StropheParams.YEAR[:-1]) - int(raw_year)))
+            return str(StropheParams.YEAR[year_index])
+        else:
+            return "NaN"
+        
+    _RHYME_POS = ["A", "B", "C", "D", "E", "F", "G", "H"]
+        
+    @staticmethod
+    def rhyme_sec(rhyme_ref, current_rhyme):
+        """Return proper rhyme indicator to given reference
+        
+        Args:
+            rhyme_ref (_type_): reference number of 'A'
+            current_rhyme (_type_): current rhyme number that needs inidcation
+            
+        Returns:
+            str: rhyme indicator character
+        """
+        
+        return "X" if current_rhyme == None or current_rhyme== -1 or rhyme_ref == None or current_rhyme < rhyme_ref or current_rhyme >= rhyme_ref + len(TextManipulation._RHYME_POS) else TextManipulation._RHYME_POS[current_rhyme - rhyme_ref]
+        
+    @staticmethod
+    def __post_process_rhyme(rhyme_str: str):
+        # First Pass
+        marker_count = {marker: rhyme_str.count(marker) for marker in TextManipulation._RHYME_POS}
+        for key, val in marker_count.items():
+            # Replace all, that ocurr only once with X
+            if val == 1:
+                rhyme_str = re.sub(key, 'X', rhyme_str)
+        # Downscale higher to lower if lower not present
+        marker_count = {marker: rhyme_str.count(marker) for marker in TextManipulation._RHYME_POS}
+        for key, val in marker_count.items():
+            if val > 1 and key != 'X':
+                key_index = TextManipulation._RHYME_POS.index(key)
+                replacements = {marker: rhyme_str.count(marker) for marker in TextManipulation._RHYME_POS[:key_index]}
+                for rep_key, rep_val in replacements.items():
+                    if rep_val ==0:
+                        rhyme_str = re.sub(key, rep_key, rhyme_str)
+                        break
+                    
+        # Pass to swap letters
+        marker_index = {marker: rhyme_str.find(marker) for marker in TextManipulation._RHYME_POS if rhyme_str.find(marker) != -1}
+        keys_values = marker_index.items()
+        keys = [v[0] for v in keys_values]
+        values = [v[1] for v in keys_values]
+       
+        i = 0
+        while i < len(keys):
+            j= 0
+            while j< len(keys):
+                if TextManipulation._RHYME_POS.index(keys[j]) > TextManipulation._RHYME_POS.index(keys[i]) and values[j] < values[i]:
+                    # Swap the positions
+                    rhyme_str = re.sub(keys[j], 'Z', rhyme_str)
+                    rhyme_str = re.sub(keys[i], keys[j], rhyme_str)
+                    rhyme_str = re.sub('Z', keys[i], rhyme_str)
+                    # Need to update the value
+                    temp = values[i]
+                    values[i]= values[j]
+                    values[j] = temp
+                j+=1
+            i+=1
+        
+            
+        return rhyme_str
+            
+        
+    @staticmethod
+    def _rhyme_string(curr_rhyme_list):
+        """Translate rhyme as list of rhyming number to rhyme schema
+
+        Args:
+            curr_rhyme_list (list): Current rhyme as list of ints indicating rhyming verses
+
+        Returns:
+            str: Rhyme schema
+        """
+        rhyme_list = curr_rhyme_list.copy()
+        reference = None
+        # Give None a blank -1 rhyme id
+        for i in range(len(rhyme_list)):
+            if rhyme_list[i] != None and reference == None:
+                reference = rhyme_list[i]
+            elif rhyme_list[i] != None and rhyme_list[i] < reference:
+                reference = rhyme_list[i]
+            elif rhyme_list[i] == None:
+                 rhyme_list[i] = -1
+        
+        # With more robust post processing, this is may not needed
+               
+        # if there is valid rhyme, normalize 
+        if reference != None:
+            # sort the rhyme and get index of reference number
+            cheat_sheet =  sorted(list(set(rhyme_list[:])))
+            ref_index = cheat_sheet.index(reference)
+            # normalize the rest around this reference
+            for i in range(len(rhyme_list)):
+                idx = cheat_sheet.index(rhyme_list[i])
+                rhyme_list[i] = reference + (idx - ref_index)
+        
+                
+        rhyme_str = ""
+        for num in rhyme_list:
+           rhyme_str += TextManipulation.rhyme_sec(reference, num)
+        
+        return TextManipulation.__post_process_rhyme(rhyme_str)
+
+class TextAnalysis:
+    """Static class with methods of analysis of strings
+
+    Returns:
+        Union[str, bool, dict, numpy.ndarray]: Analyzed input
+    """
+    
+    # Possible Keys if returned type is dict
+    POET_PARAM_LIST = ["RHYME", "YEAR", "METER", "LENGTH", "END", "TRUE_LENGTH", "TRUE_END"]
+    
+    @staticmethod
+    def _is_meter(meter:str):
+        """Return if string is meter type
+
+        Args:
+            meter (str): string to analyze
+
+        Returns:
+            bool: If string is meter type
+        """
+        return meter in StropheParams.METER[:-1]
+    
+    @staticmethod
+    def _is_year(year:str):
+        """Return if string is year or special NaN
+
+        Args:
+            year (str): string to analyze
+
+        Returns:
+            bool: If string is year or special NaN
+        """
+        return (year.isdecimal() and int(year) > 1_000 and int(year) < 10_000) or year == "NaN"
+    
+    @staticmethod
+    def _rhyme_like(rhyme:str):
+        """Return if string is structured like rhyme schema
+
+        Args:
+            rhyme (str): string to analyze
+
+        Returns:
+            bool: If string is structured like rhyme schema
+        """
+        return (rhyme.isupper() and len(rhyme) >= 3 and len(rhyme) <= 6)
+    
+    @staticmethod
+    def _rhyme_vector(rhyme:str) -> np.ndarray:
+        """Create One-hot encoded rhyme schema vector from given string
+
+        Args:
+            rhyme (str): string to construct vector from 
+
+        Returns:
+            numpy.ndarray: One-hot encoded rhyme schema vector
+        """
+        
+        rhyme_vec = np.zeros(len(StropheParams.RHYME))
+        if rhyme in StropheParams.RHYME:
+            rhyme_vec[StropheParams.RHYME.index(rhyme)] = 1
+        else:
+            rhyme_vec[-1] = 1
+            
+        return rhyme_vec
+    
+    
+    @staticmethod
+    def _publish_year_vector(year_string):
+        """Construct vector of year of publishing, weighting by distance
+
+        Args:
+            year_string (str): String with publish year
+
+        Returns:
+            numpy.ndarray: Vector of bucketized One-hot encoded publish year
+        """
+        publish_year = None if not year_string.isdigit() else int(year_string)
+        publish_vector = np.zeros(len(StropheParams.YEAR))
+        if publish_year == None:
+            publish_vector[-1] = 1
+        else:
+            # Distance Part
+            #distance_weighting = [1/(1 + abs(year - publish_year)) for year in POET_YEARS_BUCKETS[:-1]] + [0]
+            #publish_vector = np.asarray(distance_weighting)
+            # Correct class correction
+            publish_vector[np.argmin( abs(np.asarray(StropheParams.YEAR[:-1]) - publish_year))] += 1
+            # Normalize
+            #publish_vector = publish_vector/np.sum(publish_vector)
+        return publish_vector  
+    
+    @staticmethod
+    def _rhyme_or_not(rhyme_str:str) -> np.ndarray:
+        """Create vector if given rhyme string is in our list of rhyme schemas
+
+        Args:
+            rhyme_str (str): string to construct vector from 
+
+        Returns:
+            numpy.ndarray: Boolean flag vector
+        """
+        rhyme_vector = np.zeros(2)
+        if rhyme_str in StropheParams.RHYME:
+            rhyme_vector[0] = 1
+        else:
+            rhyme_vector[1] = 1
+        return rhyme_vector
+        
+    @staticmethod
+    def _metre_vector(metre: str) -> np.ndarray:
+        """Create One-hot encoded metre vector from given string
+
+        Args:
+            metre (str): string to construct vector from 
+
+        Returns:
+            numpy.ndarray: One-hot encoded metre vector
+        """
+        metre_vec = np.zeros(len(StropheParams.METER))
+        if metre in StropheParams.METER:
+            metre_vec[StropheParams.METER.index(metre)] = 1
+        else:
+            metre_vec[-1] = 1            
+        return metre_vec
+    
+    @staticmethod
+    def _first_line_analysis(text:str):
+        """Analysis of parameter line for RHYME, METER, YEAR
+
+        Args:
+            text (str): parameter line string
+
+        Returns:
+            dict: Dictionary with analysis result
+        """
+        line_striped = text.strip()
+        if not line_striped:
+            return {}
+        poet_params = {}
+        # Look for each possible parameter
+        for param in line_striped.split():
+            if TextAnalysis._is_year(param):
+                # Year is Bucketized so to fit
+                poet_params["YEAR"] = TextManipulation._year_bucketor(param) 
+            elif TextAnalysis._rhyme_like(param):
+                poet_params["RHYME"] = param
+            elif TextAnalysis._is_meter(param):
+                poet_params["STROPHE_METER"] = param
+        return poet_params
+    
+    @staticmethod
+    def _is_line_length(length:str):
+        """Return if string is number of syllables parameter
+
+        Args:
+            length (str): string to analyze
+
+        Returns:
+            bool: If string is number of syllables parameter
+        """
+        return length.isdigit() and int(length) > 1 and int(length) < 100
+    
+    @staticmethod
+    def _is_line_end(end:str):
+        """Return if string is valid ending syllable/sequence parameter
+
+        Args:
+            end (str): string to analyze
+
+        Returns:
+            bool: If string is valid ending syllable/sequence parameter
+        """
+        return end.isalpha() and end.islower() and len(end) <= 5
+    
+    @staticmethod
+    def _continuos_line_analysis(text:str):
+        """Analysis of Content lines for LENGTH, TRUE_LENGTH, END, TRUE_END
+
+        Args:
+            text (str): content line to analyze
+
+        Returns:
+            dict: Dictionary with analysis result
+        """
+        # Strip line of most separators and look if its empty
+        line_striped = TextManipulation._remove_most_nonchar(text, lower_case=False).strip()
+        if not line_striped:
+            return {}
+        line_params = {}
+        # OLD MODEL
+        if text.count('#') == 0: # BASIC
+            pass
+        else:
+            for param_group in text.split('#')[:-1]:
+                for param in param_group.split():
+                    if TextAnalysis._is_meter(param.strip()):
+                        line_params["METER"] = param.strip()
+                    elif TextAnalysis._is_line_length(param.strip()):
+                        line_params["LENGTH"] = int(param.strip())
+                    elif TextAnalysis._is_line_end(param.strip()):
+                        line_params["END"] = param.strip()
+                    
+                        
+        line_params["TRUE_LENGTH"] = len(SyllableMaker.syllabify(line_striped.split('#')[-1]))     
+        line_only_char = TextManipulation._remove_all_nonchar(line_striped).strip()
+        if len(line_only_char) > 2:
+            line_params["TRUE_END"] = SyllableMaker.syllabify(" ".join(line_only_char.split()[-2:]))[-1]
+        
+        return line_params
+    
+    @staticmethod
+    def _is_param_line(text:str):
+        """Return if line is a Parameter line (Parameters RHYME, METER, YEAR)
+
+        Args:
+            text (str): line to analyze
+
+        Returns:
+            bool: If line is a Parameter line
+        """
+        line_striped = text.strip()
+        if not line_striped:
+            return False
+        small_analysis = TextAnalysis._first_line_analysis(line_striped)
+        return  "RHYME" in small_analysis.keys() or "YEAR" in small_analysis.keys()
+
+class SyllableMaker:
+    """Static class with methods for separating string to list of Syllables
+
+    Returns:
+        list: List of syllables
+    """
+    
+        
+# NON-Original code!
+# Taken from Barbora Štěpánková
+
+    @staticmethod
+    def syllabify(text : str) -> list[str]:
+        words = re.findall(r"[aábcčdďeéěfghiíjklmnňoópqrřsštťuúůvwxyýzžAÁBCČDĎEÉĚFGHIÍJKLMNŇOÓPQRŘSŠTŤUÚŮVWXYÝZŽäöüÄÜÖ]+", text)
+        syllables : list[str] = []
+
+        i = 0
+        while i < len(words):
+            word = words[i]
+
+            if (word.lower() == "k" or word.lower() == "v" or word.lower() == "s" or word.lower() == "z") and i < len(words) - 1 and len(words[i + 1]) > 1:
+                i += 1
+                word = word + words[i]
+
+            letter_counter = 0
+
+            # Get syllables: mask the word and split the mask
+            for syllable_mask in SyllableMaker.__split_mask(SyllableMaker.__create_word_mask(word)):
+                word_syllable = ""
+                for character in syllable_mask:
+                    word_syllable += word[letter_counter]
+                    letter_counter += 1
+
+                syllables.append(word_syllable)
+
+            i += 1
+
+        return syllables
+
+
+    @staticmethod
+    def __create_word_mask(word : str) -> str:
+        word = word.lower()
+
+        vocals = r"[aeiyouáéěíýóůúäöü]"
+        consonants = r"[bcčdďfghjklmnňpqrřsštťvwxzž]"
+
+        replacements = [
+            #double letters
+    		('ch', 'c0'),
+    		('rr', 'r0'),
+            ('ll', 'l0'),
+    		('nn', 'n0'),
+    		('th', 't0'),
+
+            # au, ou, ai, oi
+    		(r'[ao]u', '0V'),
+            (r'[ao]i','0V'),
+
+            # eu at the beginning of the word
+    		(r'^eu', '0V'),
+    
+            # now all vocals
+    		(vocals, 'V'),
+
+            # r,l that act like vocals in syllables
+    		(r'([^V])([rl])(0*[^0Vrl]|$)', r'\1V\3'),
+
+            # sp, st, sk, št, Cř, Cl, Cr, Cv
+    		(r's[pt]', 's0'),
+    		(r'([^V0lr]0*)[řlrv]', r'\g<1>0'),
+    		(r'([^V0]0*)sk', r'\1s0'),
+    		(r'([^V0]0*)št', r'\1š0'),
+
+    		(consonants, 'K')
+    	]
+
+        for (original, replacement) in replacements:
+            word = re.sub(original, replacement, word)
+
+        return word
+
+
+    @staticmethod
+    def __split_mask(mask : str) -> list[str]:
+        replacements = [
+    		# vocal at the beginning
+    		(r'(^0*V)(K0*V)', r'\1/\2'),
+    		(r'(^0*V0*K0*)K', r'\1/K'),
+
+    		# dividing the middle of the word
+    		(r'(K0*V(K0*$)?)', r'\1/'),
+    		(r'/(K0*)K', r'\1/K'),
+    		(r'/(0*V)(0*K0*V)', r'/\1/\2'),
+    		(r'/(0*V0*K0*)K', r'/\1/K'),
+
+    		# add the last consonant to the previous syllable
+    		(r'/(K0*)$', r'\1/')
+    	]
+
+        for (original, replacement) in replacements:
+            mask = re.sub(original, replacement, mask)
+
+        if len(mask) > 0 and mask[-1] == "/":
+            mask = mask[0:-1]
+
+        return mask.split("/")

utils/validators.py

@@ -0,0 +1,359 @@
+import torch
+import transformers
+import jellyfish
+from tqdm import tqdm
+from transformers import  AutoModelForMaskedLM
+from transformers.utils import ModelOutput
+import numpy as np
+from .poet_utils import StropheParams
+
+from torch.utils.data import DataLoader, Dataset
+from pytorch_optimizer import SAM
+
+class ValidatorInterface(torch.nn.Module):
+    """Pytorch Model Interface. Abstract class for all validators
+
+    Args:
+        torch (_type_): Is child of torch.nn.Module for integration with torch and huggingface 
+    """
+    def __init__(self, *args, **kwargs) -> None:
+        """ Constructor. As child Class needs to construct Parent
+        """
+        super().__init__(*args, **kwargs)
+        
+    def forward(self, input_ids=None, attention_mask=None, *args, **kwargs):
+        """Compute model output and model loss
+
+        Args:
+            input_ids (_type_, optional): Model inputs. Defaults to None.
+            attention_mask (_type_, optional): Attention mask where padding starts. Defaults to None.
+
+        Raises:
+            NotImplementedError: Abstract class
+        """
+        raise NotImplementedError()
+    
+    def predict_state(self, input_ids=None, *args, **kwargs):
+        """Compute model outputs
+
+        Args:
+            input_ids (_type_, optional): Model inputs. Defaults to None.
+
+        Raises:
+            NotImplementedError: Abstract class
+        """
+        raise NotImplementedError()
+    
+    def validate_model(self, input_ids=None, *args, **kwargs):
+        """Validate model given some labels, Doesn't use loss
+
+        Args:
+            input_ids (_type_, optional): Model inputs. Defaults to None.
+
+        Raises:
+            NotImplementedError: Abstract class
+        """
+        raise NotImplementedError()
+    
+    
+class RhymeValidator(ValidatorInterface):
+    def __init__(self, pretrained_model, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        
+        self.model = AutoModelForMaskedLM.from_pretrained(pretrained_model, output_hidden_states=True)
+        
+        self.config = self.model.config
+        
+        self.model_size = self.config.hidden_size 
+        
+        self.rhyme_regressor = torch.nn.Linear(self.model_size, len(StropheParams.RHYME)) # Common Rhyme Type
+        
+        self.loss_fnc = torch.nn.CrossEntropyLoss(label_smoothing=0.0, weight=torch.tensor([1, 1, 1.5, 1.5, 1.5, 1.5, 
+                                                                                 2, 2,   2,   3,   3,   3, 
+                                                                                 3, 3,   3,   3,   4,   4, 
+                                                                                 5, 5,   5,   5,   7,   7, 
+                                                                                 7, 7,   7,   8,   8,   8,
+                                                                                 9, 9,   9,  10,  10,  10,
+                                                                                 12,12, 12,  12,  12,  12,
+                                                                                 15,15,1.5]) )
+                                                        
+    def forward(self, input_ids=None, attention_mask=None, rhyme=None, *args, **kwargs):
+        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids.type(torch.LongTensor))
+        
+        last_hidden = outputs['hidden_states'][-1]
+        
+        rhyme_regression = self.rhyme_regressor((last_hidden[:,0,:].view(-1, self.model_size)))
+            
+        softmaxed = torch.softmax(rhyme_regression, dim=1)
+        rhyme_loss = self.loss_fnc(softmaxed, rhyme)
+        
+        return ModelOutput(loss=rhyme_loss + outputs.loss, model_output=softmaxed)
+        
+    def predict_state(self, input_ids=None, *args, **kwargs):
+        
+        outputs = self.model(input_ids=input_ids)
+        
+        last_hidden = outputs['hidden_states'][-1]
+        
+        rhyme_regression = self.rhyme_regressor((last_hidden[:,0,:].view(-1, self.model_size)))
+            
+        softmaxed = torch.softmax(rhyme_regression, dim=1)
+        
+        return softmaxed
+    
+    def validate_model(self, input_ids=None, rhyme=None, k:int = 2,*args, **kwargs):
+        outputs = self.model(input_ids=input_ids)
+        
+        last_hidden = outputs['hidden_states'][-1]
+        
+        rhyme_regression = self.rhyme_regressor((last_hidden[:,0,:].view(-1, self.model_size)))
+            
+        softmaxed = torch.softmax(rhyme_regression, dim=1)
+        
+        softmaxed = softmaxed.flatten().cpu()
+        
+        predicted_val = torch.argmax(softmaxed)
+        
+        predicted_top_k = torch.topk(softmaxed, k).indices
+        
+        label_val = torch.argmax(rhyme.flatten())
+        
+        validation_true_val = (label_val == predicted_val).float().sum().numpy()
+        top_k_presence = 0
+        if label_val in predicted_top_k:
+            top_k_presence = 1
+            
+        levenshtein = jellyfish.levenshtein_distance(StropheParams.RHYME[predicted_val] if StropheParams.RHYME[predicted_val] != None else "", StropheParams.RHYME[label_val] if  StropheParams.RHYME[label_val] != None else "")
+        
+        hit_pred = softmaxed[label_val].detach().numpy()
+        
+        return {"acc" : validation_true_val,
+                "top_k" : top_k_presence,
+                "lev_distance": levenshtein,
+                "predicted_label" : hit_pred
+        }
+         
+    
+    
+class MeterValidator(ValidatorInterface):
+    def __init__(self, pretrained_model, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.model = AutoModelForMaskedLM.from_pretrained(pretrained_model, output_hidden_states=True)
+        
+        self.config = self.model.config
+        
+        self.model_size = self.config.hidden_size
+        
+        self.meter_regressor = torch.nn.Linear(self.model_size, len(StropheParams.METER)) # Meter Type
+        
+        self.loss_fnc = torch.nn.CrossEntropyLoss(label_smoothing=0.0, weight=torch.tensor([1, 1.5, 5, 10, 10, 20, 5, 20, 20, 0]))
+        
+    def forward(self, input_ids=None, attention_mask=None, metre_ids=None, *args, **kwargs):
+        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids.type(torch.LongTensor))
+        
+        last_hidden = outputs['hidden_states'][-1]
+        
+        meter_regression = self.meter_regressor((last_hidden[:,0,:].view(-1, self.model_size)))
+            
+        softmaxed = torch.softmax(meter_regression, dim=1)
+        meter_loss = self.loss_fnc(softmaxed, metre_ids)
+        
+        return ModelOutput(loss=meter_loss + outputs.loss, model_output=softmaxed)
+        
+    def predict_state(self, input_ids=None, *args, **kwargs):
+        outputs = self.model(input_ids=input_ids)
+        
+        last_hidden = outputs['hidden_states'][-1]
+        
+        meter_regression = self.meter_regressor((last_hidden[:,0,:].view(-1, self.model_size)))
+            
+        softmaxed = torch.softmax(meter_regression, dim=1)
+        
+        return softmaxed
+    
+    def validate_model(self, input_ids=None, metre_ids=None, attention_mask=None, k: int=2,*args, **kwargs):
+        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask )
+        
+        last_hidden = outputs['hidden_states'][-1]
+        
+        meter_regression = self.meter_regressor((last_hidden[:,0,:].view(-1, self.model_size)))
+            
+        softmaxed = torch.softmax(meter_regression, dim=1)
+        
+        softmaxed = softmaxed.flatten().cpu()
+        
+        predicted_val = torch.argmax(softmaxed)
+        
+        predicted_top_k = torch.topk(softmaxed, k).indices
+        
+        label_val = torch.argmax(metre_ids.flatten())
+        
+        validation_true_val = (label_val == predicted_val).float().sum().numpy()
+        top_k_presence = 0
+        if label_val in predicted_top_k:
+            top_k_presence = 1
+        
+        hit_pred = softmaxed[label_val].detach().numpy()
+        
+        return {"acc" : validation_true_val,
+                "top_k" : top_k_presence,
+                "predicted_label" : hit_pred
+        }
+        
+class YearValidator(ValidatorInterface):
+    def __init__(self, pretrained_model, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.model = AutoModelForMaskedLM.from_pretrained(pretrained_model, output_hidden_states=True)
+        
+        self.config = self.model.config
+        
+        self.model_size = self.config.hidden_size
+        
+        self.year_era = torch.nn.Linear(self.model_size, len(StropheParams.YEAR))
+        self.softmax = torch.nn.Softmax(dim=-1)
+        
+        self.year_val = torch.nn.Linear(self.model_size, 1) # Year Value     
+        
+        
+        self.loss_fnc_era = torch.nn.CrossEntropyLoss(label_smoothing=0.0,weight=torch.tensor([10, 5, 3, 3, 1, 1, 1.5, 2, 5, 0]))
+        
+        self.loss_fnc_val = torch.nn.L1Loss()
+        
+    def forward(self, input_ids=None, attention_mask=None, year_bucket=None, year=None, *args, **kwargs):
+        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids.type(torch.LongTensor))
+        
+        last_hidden = outputs['hidden_states'][-1]
+        
+        
+        year_val = self.year_val((last_hidden[:,0,:].view(-1, self.model_size)))
+        year_val_loss = self.loss_fnc_val(year_val, year)
+        
+        year_era = self.year_era((last_hidden[:,0,:].view(-1, self.model_size)))
+        year_era = self.softmax(year_era)
+        year_era_loss =  self.loss_fnc_era(year_era, year_bucket)
+        
+        return ModelOutput(loss=year_val_loss + year_era_loss  + outputs.loss, model_output=(year_val, year_era))
+        
+    def predict_state(self, input_ids=None, *args, **kwargs):
+        outputs = self.model(input_ids=input_ids)
+        
+        last_hidden = outputs['hidden_states'][-1]
+        
+        year_val = self.year_val((last_hidden[:,0,:].view(-1, self.model_size)))
+        
+        return year_val
+    
+    def validate_model(self, input_ids=None, year_bucket=None, k: int=2,*args, **kwargs):
+        
+        outputs = self.model(input_ids=input_ids)
+        
+        last_hidden = outputs['hidden_states'][-1]
+        
+        year_val = self.year_val((last_hidden[:,0,:].view(-1, self.model_size)))
+        if hasattr(self, 'year_era'):
+            year_era = self.year_era((last_hidden[:,0,:].view(-1, self.model_size)))
+            year_era = self.softmax(year_era)
+        
+        year_val = year_val.detach().flatten().cpu().numpy()
+        if hasattr(self, 'year_era'):
+            year_era = year_era.detach().flatten().cpu().numpy()
+        
+        publish_vector  = [1/(1 + abs(year - year_val[0])) for year in StropheParams.YEAR[:-1]] + [0]
+        publish_vector = np.asarray(publish_vector)/np.sum(publish_vector)
+        # Adding era prediction
+        if hasattr(self, 'year_era'):
+            publish_vector+= year_era
+        publish_vector = torch.tensor( np.asarray(publish_vector)/np.sum(publish_vector))
+        
+        
+        predicted_val = torch.argmax(publish_vector)
+        
+        predicted_top_k = torch.topk(publish_vector, k).indices
+        
+        label_val = torch.argmax(year_bucket.flatten())
+        
+        validation_true_val = (label_val == predicted_val).float().sum().numpy()
+        top_k_presence = 0
+        if label_val in predicted_top_k:
+            top_k_presence = 1
+        
+        hit_pred = publish_vector[label_val].detach().numpy()
+        
+        distance = abs(label_val.numpy() - predicted_val.numpy())
+        
+        return {"acc" : validation_true_val,
+                "top_k" : top_k_presence,
+                "predicted_label" : hit_pred,
+                "distance" : distance
+        }       
+    
+    
+
+class ValidatorTrainer:
+    def __init__(self, model: ValidatorInterface, args: dict, train_dataset: Dataset, data_collator, device):
+        self.model = model
+        self.args = args
+        self.epochs = 1 if "epochs" not in args.keys() else args["epochs"]
+        self.batch_size = 1 if "batch_size" not in args.keys() else args["batch_size"]
+        self.lr = 5e-5 if "lr" not in args.keys() else args["lr"]
+        self.weight_decay = 0.0 if "weight_decay" not in args.keys() else args['weight_decay']
+        
+        self.train_loader = DataLoader(train_dataset, self.batch_size, True, collate_fn=data_collator)
+        
+        # SAM Values
+        self.device = device      
+        self.optimizer = SAM(self.model.parameters(), torch.optim.AdamW, lr=self.lr, weight_decay=self.weight_decay)
+        self.scheduler = transformers.get_constant_schedule_with_warmup(self.optimizer, 4 * len(train_dataset)//self.batch_size)
+        
+        # GSAM Value
+        #self.device = device
+        #self.base_optim =  AdamP(self.model.parameters(), lr=self.lr, weight_decay=self.weight_decay)
+        #self.scheduler = transformers.get_constant_schedule_with_warmup(self.base_optim, len(train_dataset)//self.batch_size)
+        #self.rho_scheduler=  ProportionScheduler( self.scheduler, max_lr=self.lr)
+        #self.optimizer = GSAM(self.model.parameters(),self.base_optim, self.model, self.rho_scheduler, alpha=0.05)
+      
+    def train(self):
+        for epoch in  tqdm(range(self.epochs)):
+            self.model.train()
+            
+            # SAM Attempt
+            
+            for step, batch in enumerate(self.train_loader):
+                # First Pass
+                loss = self.model(input_ids=batch["input_ids"].to(self.device), attention_mask=batch["attention_mask"].to(self.device),
+                                  rhyme = None if batch["rhyme"] == None else batch["rhyme"].to(self.device),
+                                  metre_ids = None if batch["metre_ids"] == None else batch["metre_ids"].to(self.device),
+                                  year_bucket = None if batch["year_bucket"] == None else batch["year_bucket"].to(self.device),
+                                  year = None if batch["year"] == None else batch["year"].to(self.device))['loss']
+                loss.backward()          
+                self.optimizer.first_step(zero_grad=True)
+                # Second Pass
+                loss = self.model(input_ids=batch["input_ids"].to(self.device), attention_mask=batch["attention_mask"].to(self.device),
+                                      rhyme = None if batch["rhyme"] == None else batch["rhyme"].to(self.device),
+                                      metre_ids = None if batch["metre_ids"] == None else batch["metre_ids"].to(self.device),
+                                      year_bucket = None if batch["year_bucket"] == None else batch["year_bucket"].to(self.device),
+                                      year = None if batch["year"] == None else batch["year"].to(self.device))['loss']
+                
+                loss.backward()
+                self.optimizer.second_step(zero_grad=True)
+                self.scheduler.step()
+           
+            # GSAM Attempt 
+                 
+            #for step, batch in enumerate(self.train_loader):
+            #    def closure():
+            #        self.optimizer.base_optimizer.zero_grad()
+            #        with torch.enable_grad():
+            #            outputs = self.model(input_ids=batch["input_ids"].to(self.device), attention_mask=batch["attention_mask"].to(self.device),
+            #                      rhyme = None if batch["rhyme"] == None else batch["rhyme"].to(self.device),
+            #                      metre = None if batch["metre"] == None else batch["metre"].to(self.device))
+            #            loss = torch.nn.functional.cross_entropy(outputs['model_output'].to(self.device),batch['rhyme'].to(self.device) if isinstance(self.model, RhymeValidator) else batch['metre'].to(self.device))
+            #        loss.backward()
+            #        return outputs['model_output'], loss.detach()
+            #    predictions, loss = self.optimizer.step(closure)
+            #    self.scheduler.step()
+            #    self.optimizer.update_rho_t()
+            #    
+                if step % 100 == 0:
+                    print(f'Step {len(self.train_loader) * epoch + step},  loss : {loss.item()}', flush=True)
+    

utils/validators/meter/ufal-robeczech-base_BPE_validator_1704126400265

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d83f2b8f9b00db0945584e3bcbce96f971cfc572cb8665ff713c6d3cc67854d4
+size 504173324

utils/validators/rhyme/distilroberta-base_BPE_validator_1704126399565

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ceb77ef356a5e5ce3d59a6b2d31b96c925af09e29b4731c143ebabdaf3401c65
+size 328898329

utils/validators/year/ufal-robeczech-base_BPE_validator_1702393305267

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4695ae160b8236b89c467fb50318c6cb429ae6152f9332f74ddcaff5cbe23da1
+size 504177816