<jupyter_start><jupyter_text>Setup<jupyter_code>cd /home/jovyan/assets
# -- Base -- #
import os
import random
import joblib
import logging
import time
import re
from copy import deepcopy
from dataclasses import dataclass
import sys
import yaml
import csv
from typing import (
    List,
    Dict,
    Tuple
)
from yaspin import yaspin
from functools import partial

# -- Tokenizer -- #
import tokenizers
from tokenizers.models import WordPiece

from tokenizers import (
    Tokenizer,
    normalizers
)

from tokenizers.normalizers import (
    Lowercase,
    NFD,
    StripAccents
)

from tokenizers.pre_tokenizers import Whitespace
from tokenizers.processors import TemplateProcessing
from tokenizers.trainers import WordPieceTrainer
from tokenizers import decoders

# -- PreTrained BERT -- #
from transformers import create_optimizer
from transformers import PreTrainedTokenizerFast
from datasets import Dataset

# -- Metrics -- #
import numpy as np
import pandas as pd
import sqlite3 as sql

# -- Tensorflow -- #
import tensorflow as tf

# -- Misc Models -- #
import drain3
from gensim.models.phrases import Phrases
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import train_test_split

# -- Custom -- #
from libs.transformers.src.transformers.models.bert.modeling_tf_bert import TFBertForPreTraining
from libs.transformers.src.transformers.models.bert.configuration_bert import BertConfig
from libs.transformers.src.transformers.modeling_tf_utils import shape_list
from libs.transformers.src.transformers.data.data_tf_collator import TFDataCollatorForLanguageModeling

tf.__version__
def convert_sqlite_to_csv(inputFolder, ext, tableName):
    """ inputFolder - Folder where sqlite files are located. 
        ext - Extension of your sqlite file (eg. db, sqlite, sqlite3 etc.)
        tableName - table name from which you want to select the data.
    """
    csvWriter = csv.writer(open(inputFolder+'/output.csv', 'w', newline=''))
    for file1 in os.listdir(inputFolder):
        if file1.endswith('.'+ext):
            conn = sql.connect(inputFolder+'/'+file1)
            cursor = conn.cursor()
            cursor.execute("SELECT * FROM "+tableName)
            rows = cursor.fetchall()
            for row in rows:
                csvWriter.writerow(row)
            continue
        else:
            continue<jupyter_output><empty_output><jupyter_text>Extensions<jupyter_code>gpu_devices = tf.config.experimental.list_physical_devices('GPU')
for device in gpu_devices:
    tf.config.experimental.set_memory_growth(device, True)
print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))<jupyter_output><empty_output><jupyter_text>Environmental Variables<jupyter_code>SOURCE = '/home/' + os.environ['USER']
CONTAINER = 'core.soaesb'<jupyter_output><empty_output><jupyter_text>Logging<jupyter_code>logging.basicConfig(
    format='%(asctime)s %(levelname)s | %(message)s',
    level=logging.INFO,
    stream=sys.stdout
)

logger = logging.getLogger(__name__)<jupyter_output><empty_output><jupyter_text>Define Dataset  Define Database Functions<jupyter_code>def database_builder(path: str) -> pd.DataFrame():
    logger.info('Building DataFrame ...')
    (_, _, files) = next(os.walk(path))
    sql_query = 'SELECT * FROM logs'
    data = []
    for f in files:
        if '.db' in f:
            conn = create_connection(path + f)
            d = pd.read_sql_query(sql_query, conn)
            data.append(d)
    logger.info('...complete!')
    return pd.concat(data)


def create_connection(path: str) -> sql.Connection:
    """
    Creates a database connection
    :param path: str
        path to database object
    :return sql.Connection
        a connection to the database
    """
    try:
        conn = sql.connect(path)
        logger.info('Connected to database ' + path)
        return conn
    except sql.Error as e:
        logger.warning(e)<jupyter_output><empty_output><jupyter_text>Define Dataset Main<jupyter_code>dataset = database_builder(SOURCE + '/data/')
container_dataset = dataset[dataset['container_name'] == CONTAINER]<jupyter_output><empty_output><jupyter_text>W2V Pipeline  Pipeline Objects  Configuration<jupyter_code>def set_attributes(self, config: dict):
    try:
        config = config[self.__class__.__name__]
    except Exception as e:
        logger.warning(e)
        logger.warning('No configuration found for ' +
                       self.__class__.__name__)

    for attr in config.keys():
        setattr(self, attr, config[attr])


@dataclass
class PreprocessingGlobalConfig:
    embed_size: int = 512
    max_vocab_size: int = 2000
    buffer_size: int = 10000
    global_training: bool = True
    path: str = '/results/'

    def load(self, config):
        set_attributes(self, config)


@dataclass
class PhraserModelConfig:
    min_count: int = 5
    threshold: float = 7
    load_model: bool = True
    save_model: bool = False
    training: bool = True
    model_name: str = 'phrase_model.joblib'

    def load(self, config):
        set_attributes(self, config)


@dataclass
class TextClusteringConfig:
    load_model: bool = True
    save_model: bool = False
    training: bool = True
    model_name: str = 'template_miner.joblib'

    def load(self, config):
        set_attributes(self, config)


class PreprocessingPipelineConfig:
    def __init__(self):
        self.PreprocessingGlobalConfig = PreprocessingGlobalConfig()
        self.PhraserModelConfig = PhraserModelConfig()
        self.TextClusteringConfig = TextClusteringConfig()

    def load(self, path):
        try:
            with open(path) as f:
                preprocessing_config = yaml.load(f, Loader=yaml.FullLoader)
        except FileNotFoundError as e:
            logger.warning(e)
            return None

        self.PreprocessingGlobalConfig.load(preprocessing_config)
        self.PhraserModelConfig.load(preprocessing_config)
        self.TextClusteringConfig.load(preprocessing_config)<jupyter_output><empty_output><jupyter_text>Tokenizer<jupyter_code>class PrimeTokenizer:
    def __init__(self, max_seq_length: int):
        self.prime_tokenizer = Tokenizer(WordPiece(unk_token="[UNK]"))
        self.prime_tokenizer.normalizer = normalizers.Sequence([NFD(), Lowercase(), StripAccents()])
        self.prime_tokenizer.pre_tokenizer = Whitespace()
        self.prime_tokenizer.decoder = decoders.WordPiece()
        self.prime_tokenizer.enable_padding(length=max_seq_length)
        self.prime_tokenizer.enable_truncation(max_seq_length)

        self.prime_tokenizer.post_processor = TemplateProcessing(
            single="[CLS] $A [SEP]",
            pair="[CLS] $A [SEP] $B:1 [SEP]:1",
            special_tokens=[
                ("[CLS]", 1),
                ("[SEP]", 2),
            ],
        )

        self.trainer = WordPieceTrainer(
            vocab_size=153411,
            special_tokens=["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]"]
        )

    def text_to_sequence(self, input_) -> List[tokenizers.Encoding]:
        if type(input_) is list:
            return self.prime_tokenizer.encode_batch(input_)
        return self.prime_tokenizer.encode(input_)

    def sequence_to_text(self, input_) -> List[str]:
        if type(input_) is list:
            return self.prime_tokenizer.decode_batch(batch)
        return self.prime_tokenizer.decode(input_)

    def train(self, data):
        log_itr = iter(data)
        self.prime_tokenizer.train_from_iterator(log_itr, self.trainer)
        self.save()

    def get_tokenizer(self) -> Tokenizer:
        return self.prime_tokenizer

    def get_vocab(self) -> Dict[str, int]:
        return self.prime_tokenizer.get_vocab()
    
    def get_vocab_size(self) -> int:
        return self.prime_tokenizer.get_vocab_size()
    
    def save(self):
        self.prime_tokenizer.save(SOURCE + "/results/prime_tokenizer.json")
        
    def load(self):
        self.prime_tokenizer = Tokenizer.from_file(SOURCE + "/results/prime_tokenizer.json")<jupyter_output><empty_output><jupyter_text>Generic Save Model<jupyter_code>def save_model(model, path):
#     if not os.path.exists(path):
#         return

    if os.path.isfile(path):
        os.remove(path)
#     elif os.path.isdir(path):
#         shutil.rmtree(path)
#         return

    joblib.dump(model, path)<jupyter_output><empty_output><jupyter_text>PhraseCaptureLayer<jupyter_code>class PhraserModel:

    def __init__(self,
                 config: PhraserModelConfig,
                 global_config: PreprocessingGlobalConfig):

        super(PhraserModel, self).__init__()
        self.min_count = config.min_count
        self.threshold = config.threshold
        self.load_model = config.load_model
        self.save_model = config.save_model
        self.path = global_config.path
        self.model_name = config.model_name
        self.training = config.training

        if self.load_model:
            self.phrase_model = joblib.load(SOURCE +
                                            self.path +
                                            self.model_name)
        else:
            self.phrase_model = Phrases(min_count=self.min_count,
                                        threshold=self.threshold)

    def __call__(self, corpus: pd.DataFrame, training=None) -> list:
        if training is None:
            training = self.training
            
        def reorganize_return(corpus_with_phrases):
            log_list = []
            for tokenized_log in corpus_with_phrases:
                log_list.append(' '.join(tokenized_log))
            return log_list

        split_corpus = [log.split(' ') for log in corpus['log']]

        corpus_with_phrases = None
        if not training:
            frozen_model = self.phrase_model.freeze()
            corpus_with_phrases = self.phrase_model.__getitem__(split_corpus)
        else:
            self.phrase_model.add_vocab(split_corpus)

            if self.save_model:
                save_model(self.phrase_model, SOURCE + self.path + self.model_name)

            corpus_with_phrases = self.phrase_model.__getitem__(split_corpus)
            
        return reorganize_return(corpus_with_phrases)<jupyter_output><empty_output><jupyter_text>TextClusteringLayer<jupyter_code>class TextClustering:

    def __init__(self,
                 config: TextClusteringConfig,
                 global_config: PreprocessingGlobalConfig):

        super(TextClustering, self).__init__()
        self.load_model = config.load_model
        self.save_model = config.save_model
        self.path = global_config.path
        self.model_name = config.model_name
        self.training = config.training

        if self.load_model is True:
            self.template_miner = joblib.load(SOURCE +
                                              self.path +
                                              self.model_name)
        else:
            self.template_miner = drain3.TemplateMiner()

    def __call__(self, corpus: list, training=None) -> list:
        if training is None:
            training = self.training
            
        if training:
            for log in corpus:
                self.template_miner.add_log_message(log)
            if self.save_model:
                save_model(self.template_miner,
                           SOURCE + self.path + self.model_name)

            for idx, log in enumerate(corpus):
                template = self.template_miner.match(log).get_template()
                corpus[idx] = template

            return [re.sub(pattern=r' +',
                           repl=' ',
                           string=cluster) for cluster in corpus]
        else:
#             log_list = self.get_unique_templates()
#             print(f'Length of the log list: {len(log_list)}')
#             return log_list
            log_list = []
            log_set = set()
            for log in corpus:
                match_cluster = self.template_miner.match(log)
                if match_cluster is None:
                    match_cluster = self.template_miner.add_log_message(log)['template_mined']
                    log_set.add(match_cluster)
                else:
                    log_set.add(match_cluster.get_template())
        
#             l = [re.sub(pattern=r' +',
#                            repl=' ',
#                            string=cluster) for cluster in log_list]
            return list(log_set)
        
    def get_unique_templates(self) -> list:
        template_list = []
        for cluster in self.template_miner.drain.clusters:
            template_list.append(cluster.get_template())
        return [re.sub(pattern=r' +',
                       repl=' ',
                       string=cluster) for cluster in template_list]<jupyter_output><empty_output><jupyter_text>Preprocessing Pipeline<jupyter_code>def process_all_batches(n_iter, log_labels, batch_size):
    batches = []

    for idx in range(n_iter + 1):
        log_batch, labels = process_batch(dataset, idx, log_labels, batch_size)

        batches.append((log_batch, labels))

    return batches

def process_batch(dataset: pd.DataFrame,
                  idx: int,
                  labels: dict,
                  batch_size: int) -> tuple:
    start_window = idx * batch_size
    end_window = (idx + 1) * batch_size
    batched_data = dataset.iloc[start_window:end_window]
    encoded_batch = prime_tokenizer.text_to_sequence(batched_data['log'].to_list())
    id_batch = [log.ids for log in encoded_batch]
    y_batch = [labels[idx] for idx in batched_data['label']]

    tf_idf = tf.convert_to_tensor(id_batch, dtype=tf.float32)
    y_idf  = tf.convert_to_tensor(y_batch, dtype=tf.float32)
    
    return tf_idf, y_idf<jupyter_output><empty_output><jupyter_text>BERT Model  Custom Callback<jupyter_code># region Helper classes
class SavePretrainedCallback(tf.keras.callbacks.Callback):
    # Hugging Face models have a save_pretrained() method that saves both the weights and the necessary
    # metadata to allow them to be loaded as a pretrained model in future. This is a simple Keras callback
    # that saves the model with this method after each epoch.
    def __init__(self, output_dir, model, **kwargs):
        super().__init__()
        self.output_dir = output_dir
        self.model = model

    def on_epoch_end(self, epoch, logs=None):
        self.model.save_pretrained(self.output_dir)<jupyter_output><empty_output><jupyter_text>Model<jupyter_code>class PhoBert:
    def __init__(self, vocab_size, batch_size, epochs, batches_per_epoch, token_len):
        self.bert_config = BertConfig(
            vocab_size=vocab_size,
            hidden_size=512,
            num_hidden_layers=8,
            num_attention_heads=8
        )

        self.optimizer, self.lr_schedule = create_optimizer(
            init_lr=1e-4,
            num_train_steps=int(5 * batches_per_epoch),
            num_warmup_steps=2,
            adam_beta1=0.9,
            adam_beta2=0.999,
            adam_epsilon=0.1,
            weight_decay_rate=0.01,
        )
        
        self.BERT = TFBertForPreTraining(self.bert_config)
        self.BERT.resize_token_embeddings(token_len)
        self.loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

        # Prepare the metrics.
        self.train_acc_metric = tf.keras.metrics.SparseCategoricalAccuracy()
        self.val_acc_metric = tf.keras.metrics.SparseCategoricalAccuracy()
        
        self.epochs = epochs
        self.batch_size = batch_size
        self.outputs = None
        
        # Save Model Callback
        self.SavePretrainedModel = SavePretrainedCallback(f"{SOURCE}/results/PreTrainedModel/", self.BERT)
        
        # Checkpoint System
        self.chkpt = tf.train.Checkpoint(step=tf.Variable(1), optimizer=self.optimizer, model=self.BERT)
        self.chkpt_manager = tf.train.CheckpointManager(self.chkpt, f"{SOURCE}/results/PreTrainedChkpts/tf_ckpts", max_to_keep=3)
        
        
        self.chkpt.restore(self.chkpt_manager.latest_checkpoint)
        if self.chkpt_manager.latest_checkpoint:
            print("\nRestored from {}".format(self.chkpt_manager.latest_checkpoint))
        else:
            print("\nInitializing from scratch.")

        
    def compute_loss(self, labels, logits):
        # make sure only labels that are not equal to -100
        # are taken into account as loss
        masked_lm_active_loss = tf.not_equal(tf.reshape(tensor=labels["labels"], shape=(-1,)), -100)
        masked_lm_reduced_logits = tf.boolean_mask(
            tensor=tf.reshape(tensor=logits[0], shape=(-1, shape_list(logits[0])[2])),
            mask=masked_lm_active_loss,
        )
        masked_lm_labels = tf.boolean_mask(
            tensor=tf.reshape(tensor=labels["labels"], shape=(-1,)), mask=masked_lm_active_loss
        )
        next_sentence_active_loss = tf.not_equal(tf.reshape(tensor=labels["next_sentence_label"], shape=(-1,)), -100)
        next_sentence_reduced_logits = tf.boolean_mask(
            tensor=tf.reshape(tensor=logits[1], shape=(-1, 2)), mask=next_sentence_active_loss
        )
        next_sentence_label = tf.boolean_mask(
            tensor=tf.reshape(tensor=labels["next_sentence_label"], shape=(-1,)), mask=next_sentence_active_loss
        )

        return (masked_lm_labels, masked_lm_reduced_logits), (next_sentence_label, next_sentence_reduced_logits)
        
        
    @tf.function
    def train_step(self, x, y):
        with tf.GradientTape() as tape:
            logits = self.BERT(x, training=True)
            mlm, nsp = self.compute_loss(y, (logits["prediction_logits"], logits["seq_relationship_logits"]))
            masked_lm_loss = self.loss_fn(mlm[0], mlm[1])
            next_sentence_loss = self.loss_fn(nsp[0], nsp[1])
            loss_value = masked_lm_loss + next_sentence_loss

        grads = tape.gradient(loss_value, self.BERT.trainable_weights)
        self.optimizer.apply_gradients(zip(grads, self.BERT.trainable_weights))

        # Update training metric.
        self.train_acc_metric.update_state(mlm[0], mlm[1])
        self.train_acc_metric.update_state(nsp[0], nsp[1])
        
        self.outputs = logits

        return loss_value
    

    @tf.function
    def test_step(self, x, y):
        val_logits = self.BERT(x, training=False)
        mlm, nsp = self.compute_loss(y, (val_logits["prediction_logits"], val_logits["seq_relationship_logits"]))
        # Update val metrics
        self.val_acc_metric.update_state(mlm[0], mlm[1])
        self.val_acc_metric.update_state(nsp[0], nsp[1])
        
        
    def print(self, log, verbose):
        if not verbose:
            return
        print(log)
        
        
    def __call__(self, train_data, test_data, verbose=True):
        # Training BERT
        with yaspin(text="", color='blue') as sp:
            for epoch in range(self.epochs):
                sp.text = f"Epochs {epoch}/{self.epochs}"
                self.print("\nStart of epoch %d" % (epoch,), verbose)
                start_time = time.time()
                
                # Iterate over the batches of the dataset.
                for step, (x_batch_train, y_batch_train) in enumerate(train_data):
                    loss_value = self.train_step(x_batch_train, y_batch_train)
                    
                    # Increment and Save checkpoint for every 10 steps
                    self.chkpt.step.assign_add(1)
                    if int(self.chkpt.step) % 200 == 0:
                        save_path = self.chkpt_manager.save()
                        print("\nSaved checkpoint for step {}: {}".format(int(self.chkpt.step), save_path))
                        print("\nloss {:1.2f}".format(loss_value))

                    # Log every 200 batches.
                    if step % 200 == 0:
                        self.print(
                            "\nTraining loss (for one batch) at step %d: %.4f"
                            % (step, float(loss_value)),
                            verbose
                        )
                        self.print("Seen so far: %d samples" % ((step + 1) * self.batch_size), verbose)

                # Display metrics at the end of each epoch.
                train_acc = self.train_acc_metric.result()
                self.print("\nTraining acc over epoch: %.4f" % (float(train_acc),), verbose)

                # Reset training metrics at the end of each epoch
                self.train_acc_metric.reset_states()

                # Run an evaluation loop at the end of each epoch.
                for x_batch_test, y_batch_test in test_data:
                    self.test_step(x_batch_test, y_batch_test)
                eval_acc = self.val_acc_metric.result()
                self.val_acc_metric.reset_states()
                self.print("\nTesting acc: %.4f" % (float(eval_acc),), verbose)
                self.print("Time taken: %.2fs" % (time.time() - start_time), verbose)
                
                self.SavePretrainedModel.on_epoch_end(epoch)
                
            sp.text = ""
            sp.ok('✔ Finished Training Epochs')<jupyter_output><empty_output><jupyter_text>Unsupervised Learning Pipeline<jupyter_code>EncodedSeq = List[int]

def normalize_logs(logs: pd.DataFrame) -> pd.DataFrame:
    # remove timestamps and double spaces
    regexp = re.compile(
        r"""
        (?:               # Match all enclosed
        \d{4}-\d{2}-\d{2} # YYYY-MM-DD
        [\sT]             # Accept either a space or T
        \d{2}:\d{2}:\d{2} # HH:MM:SS
        ([.,]\d{3}|\s)    # Accept either a space or milliseconds
        )                 # End timestamp match
        | (?:\s{2,})      # Remove double spaces   
        | [^a-zA-Z\d:]    # Clean non-alphanumeric characters
        """, re.X)        # re.X enables comments and whitespace

    c_logs = deepcopy(logs)
    c_logs.loc[:, 'log'].replace(
        to_replace=regexp, 
        value=' ', 
        regex=True,
        inplace=True
    )

    return c_logs


def extract_unique_labels(logs: pd.DataFrame) -> dict:
    # -- Labels -- #
    label_unique = logs['label'].unique()
    binary_labels = LabelEncoder().fit_transform(label_unique)

    log_labels = {}
    # TODO: This seems a bit messy, could it be cleaned up? 
    for idx, label in enumerate(label_unique):
        log_labels.update({
            label: binary_labels[idx]
        })
    return log_labels


def create_sentence_pairing(examples):
    first_seqs = []
    nsp_labels = []

    examples["log"] = [
        line for line in examples["log"] if len(line) > 0 and not line.isspace()
    ]
    
    log_list = list(examples['log'])
    for idx in range(len(log_list)):
        first_value = log_list[idx]
        if random.random() > 0.5:
            # Pair with proper following log sequence
            second_value = log_list[(idx + 1) % len(log_list)]

            # IsNext Label
            nsp_labels.append(0)
        else:
            # Pair with random log
            rand_idx = random.randint(0, len(log_list) - 1)
            second_value = log_list[rand_idx]

            # IsNotNext Label
            nsp_labels.append(1)

        first_seqs.append((first_value, second_value))

    return {"log": first_seqs, "next_sentence_label": nsp_labels}


def generate_test_train_split(tokenized_datasets, test_size=.30):
    # Train - Test Split
    train_indices, test_indices = train_test_split(
        list(range(len(tokenized_datasets))), test_size=test_size
    )

    test_dataset = tokenized_datasets.select(test_indices)
    train_dataset = tokenized_datasets.select(train_indices)

    return train_dataset, test_dataset


class UnsupervisedLearningPipeline:
    def __init__(
        self, 
        config: PreprocessingPipelineConfig, 
        epochs: int = 3, 
        batch_size: int = 50, 
        seq_length: int = 200
    ) -> None:

        # Logs
        self.normalized_logs = None
        self.log_labels = None
        self.logs_with_phrases = list()
        
        # Tokenizers
        self.bert_tokenizer = PrimeTokenizer(seq_length)
        self.fast_tokenizer = None
        
        # Dataset for Training/Evaluation
        self.data_collator = None
        self.tf_train_dataset = None
        self.tf_test_dataset = None
        self.token_logs = None
                
        # Models
        self.PhoBert = None
        
        self.pm = PhraserModel(config.PhraserModelConfig, 
                               config.PreprocessingGlobalConfig)
        self.tc = TextClustering(config.TextClusteringConfig,
                                 config.PreprocessingGlobalConfig)
        
        # Hyperparameters
        self.max_seq_len = seq_length
        self.n_logs = 0
        self.n_iter = 0
        self.epochs = epochs
        self.batch_size = batch_size
        self.batches_per_epoch = None
        self.logs_as_templates = None
        

    def initialize_fast_tokenizer(self):
        tokenizer_obj = self.bert_tokenizer.get_tokenizer()
        fast_tokenizer = PreTrainedTokenizerFast(tokenizer_object=tokenizer_obj)
        fast_tokenizer.model_max_length = self.max_seq_len
        fast_tokenizer.unk_token = "[UNK]"
        fast_tokenizer.sep_token = "[SEP]"
        fast_tokenizer.pad_token = "[PAD]"
        fast_tokenizer.cls_token = "[CLS]"
        fast_tokenizer.mask_token = "[MASK]"
        self.fast_tokenizer = fast_tokenizer
        return fast_tokenizer
    
    
    @staticmethod
    def generate_data_signatures(train_dataset, test_dataset):
        # Train Signatures
        train_signature = {
            feature: tf.TensorSpec(shape=(None,), dtype=tf.int64)
            for feature in train_dataset.features
            if feature != "special_tokens_mask" and feature != "next_sentence_label"
        }
        train_signature["next_sentence_label"] = tf.TensorSpec(shape=(), dtype=tf.int64)
        train_signature["labels"] = train_signature["input_ids"]
        train_signature = (train_signature, {"labels": train_signature["labels"], "next_sentence_label": train_signature["next_sentence_label"]})
        
        # Test Signatures
        test_signature = {
            feature: tf.TensorSpec(shape=(None,), dtype=tf.int64)
            for feature in test_dataset.features
            if feature != "special_tokens_mask" and feature != "next_sentence_label"
        }
        test_signature["next_sentence_label"] = tf.TensorSpec(shape=(), dtype=tf.int64)
        test_signature["labels"] = test_signature["input_ids"]
        test_signature = (test_signature, {"labels": test_signature["labels"], "next_sentence_label": test_signature["next_sentence_label"]})
        
        return train_signature, test_signature
    

    def tokenize_function(self, examples):
        # Remove empty lines
        return self.fast_tokenizer(
            examples["log"],
            padding=False,
            truncation=True,
            max_length=self.max_seq_len,
            return_special_tokens_mask=True,
        )


    def get_pre_training_data(self):
        fast_tokenizer = self.initialize_fast_tokenizer()
        
        dt = self.normalized_logs.drop(["label", "container_name", "timestamp"], axis=1)
        data = Dataset.from_pandas(dt)
        
        tokenized_datasets = data.map(
            create_sentence_pairing,
            batched=True,
            num_proc=1,
            remove_columns=["log"],
            desc="Creating sentence pairings for NSP Head"
        )
        tokenized_datasets = tokenized_datasets.map(
            self.tokenize_function,
            batched=True,
            num_proc=1,
            remove_columns=["log"],
            desc="Running tokenizer on dataset line_by_line"
        )
        
        train_dataset, test_dataset = generate_test_train_split(tokenized_datasets)
        train_signature, test_signature = self.generate_data_signatures(train_dataset, test_dataset)
        
        self.batches_per_epoch = len(train_dataset) // self.batch_size
        
        data_collator = TFDataCollatorForLanguageModeling(
            tokenizer=self.fast_tokenizer,
            padding_length=self.max_seq_len,
            batch_size=self.batch_size
        )
        
        tokenized_generator = partial(data_collator, train_dataset, fast_tokenizer)
        test_generator = partial(data_collator, test_dataset, fast_tokenizer)
        
        self.tf_train_dataset = (
            tf.data.Dataset.from_generator(tokenized_generator, output_signature=train_signature)
            .batch(batch_size=self.batch_size, drop_remainder=True)
            .repeat(int(5))
        )
        
        self.tf_test_dataset = (
            tf.data.Dataset.from_generator(test_generator, output_signature=test_signature)
            .batch(batch_size=self.batch_size, drop_remainder=True)
        )


    def train_bert_tokenizer(self, load_model=False):
        if load_model:
            self.bert_tokenizer.load()
        else:
            self.bert_tokenizer.train(self.logs_as_templates)
            
            
    def pre_train_bert(self):
        # PhoBert Model instantiation
        self.PhoBert = PhoBert(
            self.bert_tokenizer.get_vocab_size(),
            self.batch_size,
            self.epochs,
            self.batches_per_epoch,
            len(self.fast_tokenizer)
        )
        
        # Prefetch Data
        self.tf_train_dataset = self.tf_train_dataset.prefetch(tf.data.AUTOTUNE)
        self.tf_test_dataset = self.tf_test_dataset.prefetch(tf.data.AUTOTUNE)
        
        # Start PreTraining
        self.PhoBert(self.tf_train_dataset, self.tf_test_dataset, verbose=True)
            
    
    def run_bert(self, batch):
        return self.BERT(batch,
                         training=False,
                         output_attentions=True)
    

    def fit(self, logs: pd.DataFrame):
        assert len(logs.index) > 0, 'process received an empty dataframe!'
        
        with yaspin(text="Normalizing Logs", color='green') as sp:
            self.normalized_logs = normalize_logs(logs)
            sp.text = ""
            sp.ok('✔ Completed log normalization')
            
            sp.text = "Extracting phrases"
            self.logs_with_phrases = self.pm(self.normalized_logs)
            sp.text = ""
            sp.ok('✔ Completed phrase extraction')
            
            sp.text = "Converting to log templates"
            self.logs_as_templates = np.array(self.tc(self.logs_with_phrases))
            sp.text = ""
            sp.ok('✔ Completed log template conversion')
        
            sp.text = "Extracting Unique Labels"
            self.log_labels = extract_unique_labels(self.normalized_logs)
            sp.text = ""
            sp.ok('✔ Completed extracting unique labels')

            sp.text = "Training Tokenizer"
            self.train_bert_tokenizer()
            sp.text = ""
            sp.ok('✔ Completed training of custom tokenizer')
                        
            sp.text = "Processing training dataset"
            self.get_pre_training_data()
            sp.text = ""
            sp.ok('✔ Completed processing training dataset')
            
            sp.text = "Pretraining BERT"
            self.pre_train_bert()
            sp.text = ""
            sp.ok('✔ Completed BERT pretraining')
            
        
    def transform(self, batch: pd.DataFrame):
        x = normalize_logs(batch)
        x = self.pm(x, False)
#         x = list(self.tc(x, False))
#         x = self.serve_batch(x)
        return self.run_bert(x)<jupyter_output><empty_output><jupyter_text>W2V Pipeline Main<jupyter_code>config_path = SOURCE + '/assets/notebooks/PreprocessingConfig.yaml'
preprocessing_config = PreprocessingPipelineConfig()
preprocessing_config.load(config_path)
# -- Unsupervised Learning Pipeline -- #

'''
Input: pd.DataFrame with batch_size number of rows 
Seq: 
    Normalize 
    Phraser
    Clustering
    Extract Unique Layers
    BERT
Returns: transformers.TFBertForPreTrainingOutput
'''

# --- SUBWORD TOKENIZER --
w2vp = UnsupervisedLearningPipeline(preprocessing_config, epochs=2, batch_size=10, seq_length=300)
training_outputs = w2vp.fit(dataset[:1000])
%%capture
!pip install bertviz
from bertviz import head_view, model_view
test_data['log']
sentence = test_data.iloc[0]['log']
sentence
tokens = w2vp.bert_tokenizer.text_to_sequence(sentence).tokens
tokens
head_view(attention, tokens)
model_view(attention, tokens)
w2vp.BERT.summary()
copy_bert = deepcopy(w2vp.BERT)
copy_bert.get_output_at
w2vp.BERT.save(SOURCE + '/results/')
tf_data<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.0961756330567378, 'fraction_numerical': 0.005266542147646896, 'mean_word_length': 2.356284040694687, 'pattern_counts': {'":': 6, '<': 65, '<?xml version=': 0, '>': 82, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '16', 'dir': 'jupyter-structured-clean-dedup', 'id': '152288', 'license': 'MIT', 'n_tokens_mistral': 10059, 'n_tokens_neox': 9710, 'n_words': 1848, 'path': 'assets/notebooks/pretraining.ipynb', 'repo_name': 'whytheevanssoftware/log-analyzer', 'size': '32659'}