translate-package 0.0.1__py3-none-any.whl

translate_package/__init__.py

@@ -0,0 +1,45 @@
+# -*- coding: utf-8 -*-
+from tokenizers import (
+    decoders,
+    models,
+    normalizers,
+    pre_tokenizers,
+    processors,
+    trainers,
+    Tokenizer
+)
+from transformers import (
+    GenerationConfig,
+    TrainingArguments,
+    Trainer, AutoModelForSeq2SeqLM,
+    get_linear_schedule_with_warmup,
+    T5ForConditionalGeneration, Adafactor, BartForConditionalGeneration,
+    MT5ForConditionalGeneration, AdamWeightDecay
+)
+from wolof_translate.utils.bucket_iterator import SequenceLengthBatchSampler, BucketSampler
+from wolof_translate.utils.sent_transformers import TransformerSequences
+from wolof_translate.utils.sent_corrections import *
+from peft import LoraConfig, get_peft_model, TaskType
+from sklearn.model_selection import train_test_split
+from torch.utils.data import DataLoader, Dataset
+from nlpaug.augmenter import char as nac
+import matplotlib.pyplot as plt
+import pytorch_lightning as pl
+from functools import partial
+import sentencepiece as spm
+from math import ceil
+import pandas as pd
+import numpy as np
+import argparse
+import evaluate
+import string
+import random
+import shutil
+import wandb
+import torch
+import time
+import nltk
+import os
+
+# Désactiver le parallélisme des tokenizers
+os.environ["TOKENIZERS_PARALLELISM"] = "false"

translate_package/data/data_preparation.py

@@ -0,0 +1,424 @@
+from translate_package import (
+    pd,
+    train_test_split,
+    Dataset,
+    DataLoader,
+    plt,
+    torch,
+    SequenceLengthBatchSampler,
+    BucketSampler,
+    partial,
+    Union,
+    Callable,
+    ceil,
+    np,
+    TransformerSequences,
+    nac,
+    remove_mark_space,
+    delete_guillemet_space
+)
+
+# sentence beginning with "Mooy li ko waral ci li ñu xamle waaye itam" is too long and must removed or corrected
+
+# python translate_hyperparameter_tuning.py --model_generation "t5" --model_name "google-t5/t5-small" --tokenizer_name "sp" --use_bucketing --save_artifact
+
+def augment(examples, src_label, p_word = 0.12554160436087158, p_char = 0.8269672653838092, max_words = 21):
+    
+    examples[src_label] = TransformerSequences(nac.RandomCharAug(action = 'swap', aug_word_p = p_word, aug_char_p = p_char, aug_word_max = max_words))(examples[src_label])[0]
+    
+    return examples
+
+def augment_(examples, src_label, tgt_label):
+    
+    examples[src_label] = TransformerSequences(remove_mark_space, delete_guillemet_space)(examples[src_label])[0]
+    
+    examples[tgt_label] = TransformerSequences(remove_mark_space, delete_guillemet_space)(examples[tgt_label])[0]
+    
+    return examples
+
+def tokenize(examples, tokenizer, src_label, tgt_label, model_generation):
+    
+    if model_generation in ["t5", "mt5", "nllb"]:
+        
+        eos_token = ""
+        bos_token = ""
+    
+    else:
+        
+        eos_token = tokenizer.eos_token 
+        bos_token = tokenizer.bos_token
+    
+    examples[src_label] = bos_token + examples[src_label] + eos_token
+
+    examples[tgt_label] = bos_token + examples[tgt_label] + eos_token
+
+    examples.update({key: value[0] for key, value in tokenizer(examples[src_label], return_tensors = 'pt').items()})
+
+    examples.update({f'decoder_{key}': value[0] for key, value in tokenizer(examples[tgt_label], return_tensors = 'pt').items()})
+
+    examples['labels'] = examples['decoder_input_ids']
+
+    return examples
+
+def apply_funcs(funcs, data):
+    # Logic to apply the functions
+    for func in funcs:
+        data = func(data)
+    return data
+
+def sequences(examples, functions):
+
+    for function in functions:
+        
+        examples = function(examples)
+   
+    return examples
+
+class SentenceDataset(Dataset):
+    
+    def __init__(self, dataframe, transformers: Union[Callable, None] = None, source_column: str = 'WOLOF', target_column: str = 'FRENCH'):
+        
+        assert source_column in dataframe.columns.tolist() and target_column in dataframe.columns.tolist()
+        
+        self.source_sentences = dataframe[source_column].tolist()
+        
+        self.target_sentences = dataframe[target_column].tolist()
+        
+        self.transformers = transformers
+        
+        self.source_column = source_column
+        
+        self.target_column = target_column
+    
+    def __getitem__(self, index):
+        
+        source_sentence = self.source_sentences[index]
+        
+        target_sentence = self.target_sentences[index]
+        
+        sentences = {
+            self.source_column: source_sentence,
+            self.target_column: target_sentence
+        }
+        
+        if not self.transformers is None:
+            
+            sentences = self.transformers(sentences)
+        
+        return sentences
+    
+    def __len__(self):
+        
+        return len(self.source_sentences)
+        
+
+def load_data(src_label, tgt_label, data_path, test_size, valid_size, seed):
+    
+    # load the dataset with pandas
+    dataset_ = pd.read_csv(data_path)
+
+    # split dataset between train, validation, and test sets
+    if test_size == 1.0:
+
+        dataset = {
+            "test": partial(SentenceDataset, dataframe = dataset_, source_column = src_label, target_column = tgt_label),
+        }
+
+    else:
+
+        train, test = train_test_split(
+            dataset_, test_size=test_size + valid_size, random_state=seed
+        )
+
+        valid, test = train_test_split(
+            test, test_size=test_size / (valid_size + test_size), random_state=seed
+        )
+
+        dataset = {
+            "train": partial(SentenceDataset, dataframe = train, source_column = src_label, target_column = tgt_label),
+            "val": partial(SentenceDataset, dataframe = valid, source_column = src_label, target_column = tgt_label),
+            "test": partial(SentenceDataset, dataframe = test, source_column = src_label, target_column = tgt_label),
+        }
+
+    # The dataset actually contains 3 diff splits: train, validation, test.
+
+    return dataset
+
+def get_boundaries(dataset, sizes, min_count):
+
+    length = []
+
+    for i in range(len(dataset)):
+
+        length.append(max(len(dataset[i]["input_ids"]), len(dataset[i]["labels"])))
+
+    # Create histogram
+    hist, bins, _ = plt.hist(length, bins=10)  # Adjust the number of bins as needed
+
+    # Analyze the histogram
+    # Identify peaks or gaps to determine the boundaries
+
+    # Choose the boundaries based on the analysis
+    boundaries = (
+        [ceil(bins[0])]
+        + [ceil(bin) for bin, count in zip(bins[1:], hist) if count > min_count]
+        + [np.inf]
+    )
+
+    boundaries = boundaries[:-1]
+
+    # define batch sizes and samplers
+    batch_sizes = [
+        sizes[i] if (i + 1) < len(sizes) else sizes[-1] for i in range(len(boundaries))
+    ]
+
+    return boundaries, batch_sizes
+
+def collate_fn_trunc(batch, input_max_len, label_max_len, eos_token_id, pad_token_id, keys: list = ['input_ids', 'attention_mask', 'labels']):
+    
+    from torch.nn.utils.rnn import pad_sequence
+
+    df_dict = {key: [] for key in keys}
+
+    for b in batch:
+
+        for key in df_dict:
+
+            df_dict[key].append(b[key])
+
+    padded_sequences = {}
+
+    for key in df_dict:
+
+        max_len = label_max_len if 'decoder' in key or 'label' in key else input_max_len
+
+        padding_value = 0 if 'mask' in key else pad_token_id # must be take care
+
+        # Pad the input sequences to have the same length
+        padded_sequences[key] = pad_sequence(df_dict[key], batch_first=True, padding_value = padding_value)[:,:max_len]
+
+        # eos token if it is not the case
+        if not 'mask' in key:
+
+             padded_sequences[key][:, -1:][(padded_sequences[key][:, -1:] != eos_token_id) & (padded_sequences[key][:, -1:] != pad_token_id)] = eos_token_id
+
+    return padded_sequences
+
+# define padding collate function
+def pad_collate(batch, padding_value):
+
+    X = [b["input_ids"] for b in batch]
+    att = [b["attention_mask"] for b in batch]
+    y = [b["labels"] for b in batch]
+
+    X_ = torch.nn.utils.rnn.pad_sequence(
+        X, batch_first=True, padding_value=padding_value
+    )
+    att_ = torch.nn.utils.rnn.pad_sequence(att, batch_first=True, padding_value=0)
+    y_ = torch.nn.utils.rnn.pad_sequence(
+        y, batch_first=True, padding_value=padding_value
+    )
+
+    return {"input_ids": X_, "attention_mask": att_, "labels": y_}
+
+
+def get_loaders(
+    tokenizer,
+    model_generation,
+    src_label,
+    tgt_label,
+    sizes,
+    data_path,
+    test_size,
+    valid_size,
+    seed,
+    p_word,
+    p_char,
+    max_words,
+    count,
+    src_max_len,
+    tgt_max_len,
+    num_workers,
+    device,
+    use_bucketing,
+    use_truncation,
+    batch_size,
+):
+
+    # get dataset
+    dataset = load_data(src_label, tgt_label, data_path, test_size, valid_size, seed)
+    
+    # analysis transformations
+    
+    a_transformers = partial(sequences, 
+                functions = [
+                    partial(augment_, src_label = src_label, tgt_label = tgt_label),
+                    partial(tokenize, tokenizer = tokenizer, src_label = src_label, tgt_label = tgt_label, model_generation = model_generation)
+            ])
+    
+    # training transformations
+    t_transformers = partial(sequences, 
+                functions = [
+                    partial(augment, src_label = src_label, p_word = p_word, p_char = p_char, max_words = max_words),
+                    partial(augment_, src_label = src_label, tgt_label = tgt_label),
+                    partial(tokenize, tokenizer = tokenizer, src_label = src_label, tgt_label = tgt_label, model_generation = model_generation)
+            ])
+    
+    if use_bucketing:
+        
+        if use_truncation:
+            
+            # initialize loaders
+            train_sampler = BucketSampler(
+                dataset["train"](transformers=a_transformers),
+                batch_size=batch_size,
+                input_key="input_ids",
+                label_key="labels",
+            )
+
+            valid_sampler = BucketSampler(
+                dataset["val"](transformers=a_transformers),
+                batch_size=batch_size,
+                input_key="input_ids",
+                label_key="labels",
+            )
+
+            test_sampler = BucketSampler(
+                dataset["test"](transformers=a_transformers),
+                batch_size=batch_size,
+                input_key="input_ids",
+                label_key="labels",
+            )
+            
+            # add transformations
+            dataset = {s: dataset[s](transformers = t_transformers) if s == 'train' else dataset[s](transformers = a_transformers) for s in dataset}
+            
+            # define data loaders
+            train_loader = DataLoader(
+                dataset["train"],
+                batch_sampler=train_sampler,
+                collate_fn = partial(collate_fn_trunc, input_max_len = src_max_len, label_max_len = tgt_max_len,
+                                    eos_token_id = tokenizer.eos_token_id, pad_token_id = tokenizer.pad_token_id),
+                num_workers=num_workers,
+                pin_memory=True if device in ["cuda", "gpu"] else False,
+            )
+            valid_loader = DataLoader(
+                dataset["val"],
+                batch_sampler=valid_sampler,
+                collate_fn=partial(collate_fn_trunc, input_max_len = src_max_len, label_max_len = tgt_max_len,
+                                    eos_token_id = tokenizer.eos_token_id, pad_token_id = tokenizer.pad_token_id),
+                num_workers=num_workers,
+                pin_memory=True if device in ["cuda", "gpu"] else False,
+            )
+            test_loader = DataLoader(
+                dataset["test"],
+                batch_sampler=test_sampler,
+                collate_fn=partial(collate_fn_trunc, input_max_len = src_max_len, label_max_len = tgt_max_len,
+                                    eos_token_id = tokenizer.eos_token_id, pad_token_id = tokenizer.pad_token_id),
+                num_workers=num_workers,
+                pin_memory=True if device in ["cuda", "gpu"] else False,
+            )
+        
+        else:
+            
+            # get boundaries
+            boundaries, batch_sizes = get_boundaries(dataset['train'](transformers = a_transformers), sizes, count)
+
+            # initialize loaders
+            train_sampler = SequenceLengthBatchSampler(
+                dataset["train"](transformers=a_transformers),
+                boundaries=boundaries,
+                batch_sizes=batch_sizes,
+                input_key="input_ids",
+                label_key="labels",
+            )
+
+            valid_sampler = SequenceLengthBatchSampler(
+                dataset["val"](transformers=a_transformers),
+                boundaries=boundaries,
+                batch_sizes=batch_sizes,
+                input_key="input_ids",
+                label_key="labels",
+            )
+
+            test_sampler = SequenceLengthBatchSampler(
+                dataset["test"](transformers=a_transformers),
+                boundaries=boundaries,
+                batch_sizes=batch_sizes,
+                input_key="input_ids",
+                label_key="labels",
+            )
+            
+            # add transformations
+            dataset = {s: dataset[s](transformers = t_transformers) for s in dataset}
+            
+            # define data loaders
+            train_loader = DataLoader(
+                dataset["train"],
+                batch_sampler=train_sampler,
+                collate_fn=partial(pad_collate, padding_value=tokenizer.pad_token_id),
+                num_workers=num_workers,
+                pin_memory=True if device in ["cuda", "gpu"] else False,
+            )
+            valid_loader = DataLoader(
+                dataset["val"],
+                batch_sampler=valid_sampler,
+                collate_fn=partial(pad_collate, padding_value=tokenizer.pad_token_id),
+                num_workers=num_workers,
+                pin_memory=True if device in ["cuda", "gpu"] else False,
+            )
+            test_loader = DataLoader(
+                dataset["test"],
+                batch_sampler=test_sampler,
+                collate_fn=partial(pad_collate, padding_value=tokenizer.pad_token_id),
+                num_workers=num_workers,
+                pin_memory=True if device in ["cuda", "gpu"] else False,
+            )
+
+    else:
+        
+        # add transformations
+        dataset = {s: dataset[s](transformers = t_transformers) for s in dataset}
+
+        if "train" in dataset:
+            # define data loaders
+            train_loader = DataLoader(
+                dataset["train"],
+                batch_size=batch_size,
+                collate_fn=partial(pad_collate, padding_value=tokenizer.pad_token_id),
+                num_workers=num_workers,
+                pin_memory=True if device in ["cuda", "gpu"] else False,
+                shuffle=True,
+            )
+
+        if "val" in dataset:
+
+            valid_loader = DataLoader(
+                dataset["val"],
+                batch_size=batch_size,
+                collate_fn=partial(pad_collate, padding_value=tokenizer.pad_token_id),
+                num_workers=num_workers,
+                pin_memory=True if device in ["cuda", "gpu"] else False,
+            )
+
+        if "test" in dataset:
+
+            test_loader = DataLoader(
+                dataset["test"],
+                batch_size=batch_size,
+                collate_fn=partial(pad_collate, padding_value=tokenizer.pad_token_id),
+                num_workers=num_workers,
+                pin_memory=True if device in ["cuda", "gpu"] else False,
+            )
+
+    if "train" in dataset and "val" in dataset:
+
+        return {
+            "train_loader": train_loader,
+            "valid_loader": valid_loader,
+            "test_loader": test_loader,
+        }
+
+    else:
+
+        return {"test_loader": test_loader}

translate_package/errors/__init__.py

@@ -0,0 +1,12 @@
+class EmptyArtifactException(Exception):
+    def __init__(
+        self, error="No wandb's artifact path was provided!! Please provide one."
+    ):
+
+        print(error)
+        
+class TokenizerException(Exception):
+    
+    def __init__(self, error: str):
+        
+        print(error)

translate_package/models/gradient_observation.py

@@ -0,0 +1,9 @@
+from torch.nn import Module
+
+def get_gradients_mean(model: Module):
+    ave_grads = []
+    for name, param in model.named_parameters():
+        if param.grad is not None:
+            ave_grads.append(param.grad.abs().mean().item())
+    
+    return sum(ave_grads)/(len(ave_grads) + 1e-5)

translate_package/models/lstm.py

@@ -0,0 +1,97 @@
+import torch
+
+class LSTMOutput:
+    
+    def __init__(self, logits, loss):
+        
+        self.logits = logits
+        
+        self.loss = loss
+
+class LSTMSequenceToSequence(torch.nn.Module):
+
+  def __init__(self, tokenizer, embedding_size = 128, num_layers = 6, hidden_size = 128, dropout=0.1, bidirectional = True):
+    
+    super().__init__()
+
+    self.tokenizer = tokenizer
+    
+    self.vocab_size = self.tokenizer.vocab_size
+
+    self.embedding = torch.nn.Embedding(self.vocab_size, embedding_size)
+
+    self.encoder = torch.nn.LSTM(input_size = embedding_size, hidden_size = hidden_size, num_layers = num_layers, batch_first = True,
+                                 bidirectional = bidirectional, dropout=dropout)
+
+    self.decoder = torch.nn.LSTM(input_size = embedding_size, hidden_size = hidden_size, num_layers = num_layers, batch_first = True,
+                                 bidirectional = bidirectional, dropout=dropout)
+
+    copy = 2 if bidirectional else 1
+
+    self.decoder_output_layer = torch.nn.Linear(copy * hidden_size, self.vocab_size)
+
+    self.loss_fn = torch.nn.CrossEntropyLoss(ignore_index=tokenizer.pad_token_id)
+
+  def forward(self, input, output):
+    
+    input_embed = self.embedding(input)
+
+    state, hidden = self.encoder(input_embed)
+
+    # decal output for decoder
+    decoder_input = output[:, :-1]
+
+    decoder_input = self.embedding(decoder_input)
+
+    decoder_output, _ = self.decoder(decoder_input, hidden)
+
+    decoder_output = self.decoder_output_layer(decoder_output)
+    
+    loss = self.loss_fn(decoder_output.reshape(-1, decoder_output.shape[-1]), output[:, 1:].reshape(-1))
+
+    return LSTMOutput(decoder_output, loss)
+
+  def generate(self, input, max_new_tokens: int = 100, temperature: float = 0.0, use_sampling = False, **kwargs):
+
+    input_embed = self.embedding(input)
+
+    _, hidden = self.encoder(input_embed)
+
+    # initialize predictions
+    predictions = torch.tensor([[self.tokenizer.bos_token_id]]*input_embed.shape[0], dtype=torch.long, device=input.device)
+    
+    # variable identifying if the sequence is finished
+    finished_sequences = torch.zeros(input_embed.shape[0], dtype=torch.bool, device = input.device)
+    
+    # generate predictions
+    for i in range(max_new_tokens):
+
+      decoder_input = self.embedding(predictions)
+
+      decoder_output, hidden = self.decoder(decoder_input, hidden)
+
+      decoder_output = self.decoder_output_layer(decoder_output)
+
+      if temperature > 0.0: decoder_output = (decoder_output / temperature)
+
+      # get probs and sample the next token from a multinomial distribution
+      probs = torch.softmax(decoder_output[:, -1], dim = -1)
+
+      if use_sampling: prediction = torch.multinomial(probs, num_samples = 1)
+      else: prediction = torch.argmax(probs, dim=-1, keepdim=True)
+
+      # add new prediction
+      predictions = torch.cat((predictions, prediction), dim = -1)
+
+      # recuperate next ids
+      next_token_ids = prediction.squeeze(-1)
+      
+      finished_sequences |= (next_token_ids == self.tokenizer.eos_token_id)
+
+      if finished_sequences.all():
+        
+        break
+      
+    # return predictions
+    return predictions[:, 1:]
+

translate_package/models/machine_translation.py

@@ -0,0 +1,308 @@
+from translate_package import (
+    pl,
+    evaluate,
+    LoraConfig,
+    TaskType,
+    torch,
+    get_linear_schedule_with_warmup,
+    wandb,
+    get_peft_model,
+    T5ForConditionalGeneration,
+    MT5ForConditionalGeneration,
+    BartForConditionalGeneration,
+    AutoModelForSeq2SeqLM,
+    Adafactor,
+    AdamWeightDecay
+)
+
+from translate_package.models.gradient_observation import get_gradients_mean
+
+from translate_package.models.lstm import LSTMSequenceToSequence
+
+
+def print_number_of_trainable_model_parameters(model):
+    trainable_model_params = 0
+    all_model_params = 0
+    for _, param in model.named_parameters():
+        all_model_params += param.numel()
+        if param.requires_grad:
+            trainable_model_params += param.numel()
+    return f"trainable model parameters: {trainable_model_params}\nall model parameters: {all_model_params}\npercentage of trainable model parameters: {100 * trainable_model_params / all_model_params:.2f}%"
+
+class MachineTranslationTransformer(pl.LightningModule):
+
+    rouge = evaluate.load("rouge")
+    bleu = evaluate.load("sacrebleu")
+
+    def __init__(
+        self,
+        model_name,
+        tokenizer,
+        model_generation="t5",
+        model=None,
+        lr=1e-4,
+        weight_decay=1e-2,
+        num_warmup_steps=0,
+        num_training_steps=20000,
+        r=32,
+        lora_alpha=32,
+        lora_dropout=0.05,
+        bias="none",
+        max_new_tokens=200,
+        predict_with_generate=True,
+        num_beams=0,
+        use_peft=False,
+        embedding_size=128,
+        num_layers=6,
+        hidden_size=128,
+        dropout=0.1,
+        bidirectional=False
+    ):
+
+        super().__init__()
+
+        if model is None:
+            if model_generation in ["t5"]:
+                
+                self.original_model = T5ForConditionalGeneration.from_pretrained(
+                    model_name, torch_dtype=torch.float32
+                )
+                
+            elif model_generation in ["mt5"]:
+                
+                self.original_model = MT5ForConditionalGeneration.from_pretrained(
+                    model_name, torch_dtype=torch.float32
+                )
+                
+            elif model_generation in ["nllb"]:
+                
+                self.original_model = AutoModelForSeq2SeqLM.from_pretrained(
+                    model_name, torch_dtype=torch.float32
+                )
+                
+            elif model_generation in ["bart"]:
+                
+                self.original_model = BartForConditionalGeneration.from_pretrained(
+                    model_name, torch_dtype=torch.float32
+                )
+            
+            elif model_generation in ["lstm"]:
+                
+                self.original_model = LSTMSequenceToSequence(tokenizer, embedding_size, num_layers, hidden_size, dropout, bidirectional)
+                
+            # resize the token embeddings
+            if not model_generation in ["lstm"]: self.original_model.resize_token_embeddings(len(tokenizer))
+            
+            if use_peft and not model_generation in ["lstm"]:
+                
+                self.lora_config = LoraConfig(
+                    r=r,  # Rank
+                    lora_alpha=lora_alpha,
+                    lora_dropout=lora_dropout,
+                    bias=bias,
+                    task_type=TaskType.SEQ_2_SEQ_LM, 
+                )
+                
+                self.model = get_peft_model(self.original_model, self.lora_config)          
+            
+            else:
+                
+                self.model = self.original_model
+            
+        else:
+
+            self.model = model
+
+        print(print_number_of_trainable_model_parameters(self.model))
+
+        self.tokenizer = tokenizer
+        
+        self.lr = lr
+
+        self.weight_decay = weight_decay
+
+        self.num_warmup_steps = num_warmup_steps
+
+        self.num_training_steps = num_training_steps
+
+        self.predict_with_generate = predict_with_generate
+
+        self.max_new_tokens = max_new_tokens
+
+        self.num_beams = num_beams
+        
+        self.model_generation = model_generation
+
+        self.predictions = {
+            "Source references": [],
+            "Predictions": [],
+            "Target references": [],
+        }
+
+    def forward(self, input):
+
+        output = self.model(**input) if not self.model_generation in ["lstm"] else self.model(input['input_ids'], input['labels'])
+
+        return output.loss, output.logits
+
+    def configure_optimizers(self):
+
+        if self.model_generation in ["t5", "mt5"]:
+            
+            optimizer = Adafactor(
+                self.parameters(), lr=self.lr, weight_decay=self.weight_decay, relative_step = False,
+                warmup_init = False
+            )
+        
+        elif self.model_generation in ["bart", "mbart"]:
+            
+            optimizer = torch.optim.AdamW(
+                self.parameters(), lr=self.lr, weight_decay=self.weight_decay
+            )
+        
+        elif self.model_generation in ["lstm"]:
+            
+            optimizer = torch.optim.AdamW(
+                self.parameters(), lr=self.lr, weight_decay=self.weight_decay
+            )
+
+        if self.model_generation in ["t5", "lstm", "mt5"]:
+            
+            return [optimizer]
+
+        elif self.model_generation in ["bart"]:
+
+            scheduler = get_linear_schedule_with_warmup(
+                optimizer,
+                num_warmup_steps=self.num_warmup_steps,
+                num_training_steps=self.num_training_steps,
+            )
+            
+            return {'optimizer': optimizer, 'lr_scheduler': {"scheduler": scheduler}}
+            
+    def training_step(self, batch, batch_idx=None):
+
+        loss, y_pred = self(batch)
+
+        self.log_dict(
+            {"train_loss": loss, "global_step": float(self.global_step)},
+            prog_bar=True,
+            on_step=False,
+            on_epoch=True,
+            sync_dist=True,
+        )
+        
+        mean_grad = get_gradients_mean(self.original_model)
+
+        wandb.log({"train_loss": loss, "trainer/global_step": self.global_step, "mean_gradient": mean_grad})
+
+        return loss
+
+    def validation_step(self, batch, batch_idx=None):
+
+        loss, y_pred = self(batch)
+
+        metrics = {}
+
+        if self.predict_with_generate:
+
+            # generate predictions
+            predictions = self.model.generate(
+                input_ids=batch["input_ids"],
+                attention_mask=batch["attention_mask"],
+                max_new_tokens=self.max_new_tokens,
+            ) if not self.model_generation in ["lstm"] else self.model.generate(
+                input=batch["input_ids"],
+                max_new_tokens=self.max_new_tokens,
+            ) 
+
+            # decode the labels
+            predictions = self.tokenizer.batch_decode(
+                predictions, skip_special_tokens=True
+            )
+            labels = self.tokenizer.batch_decode(
+                batch["labels"], skip_special_tokens=True
+            )
+
+            # get bleu metric
+            bleu = self.bleu.compute(
+                predictions=predictions,
+                references=[[label.strip()] for label in labels],
+            )
+
+            metrics["bleu"] = bleu["score"]
+
+            # get rouge metrics
+            rouge = self.rouge.compute(
+                predictions=predictions, references=[label.strip() for label in labels]
+            )
+
+            metrics.update({k: v for k, v in rouge.items() if "rouge" in k})
+
+        # get the loss
+        metrics.update(
+            {"eval_loss": loss.item(), "global_step": float(self.global_step)}
+        )
+
+        self.log_dict(
+            metrics, prog_bar=True, on_step=False, on_epoch=True, sync_dist=True
+        )
+
+        metrics.update({"trainer/global_step": self.global_step})
+
+        wandb.log(metrics)
+
+        return loss
+
+    def test_step(self, batch, batch_idx):
+
+        loss, y_pred = self(batch)
+
+        references = self.tokenizer.batch_decode(
+            batch["input_ids"], skip_special_tokens=True
+        )
+
+        # generate predictions
+        predictions = self.model.generate(
+            input_ids=batch["input_ids"],
+            attention_mask=batch["attention_mask"],
+            max_new_tokens=self.max_new_tokens,
+            do_sample=self.num_beams > 0,
+            num_beams=self.num_beams
+        ) if not self.model_generation in ["lstm"] else self.model.generate(
+            input=batch["input_ids"],
+            max_new_tokens=self.max_new_tokens,
+            use_sampling=True
+        )
+
+        # decode the labels
+        predictions = self.tokenizer.batch_decode(predictions, skip_special_tokens=True)
+        labels = self.tokenizer.batch_decode(batch["labels"], skip_special_tokens=True)
+
+        self.predictions["Source references"].extend(references)
+        self.predictions["Predictions"].extend(predictions)
+        self.predictions["Target references"].extend(labels)
+
+        # get bleu metric
+        bleu = self.bleu.compute(
+            predictions=predictions, references=[[label.strip()] for label in labels]
+        )
+
+        metrics = {}
+
+        metrics["bleu"] = bleu["score"]
+
+        # get rouge metrics
+        rouge = self.rouge.compute(predictions=predictions, references=labels)
+
+        metrics.update({k: v for k, v in rouge.items() if "rouge" in k})
+
+        # get the loss
+        metrics.update(
+            {"test_loss": loss.item(), "global_step": float(self.global_step)}
+        )
+
+        # log metrics
+        self.log_dict(
+            metrics, prog_bar=True, on_step=False, on_epoch=True, sync_dist=True
+        )

translate_package/tokenization/load_tokenizer.py

@@ -0,0 +1,40 @@
+from transformers import BartTokenizerFast
+from transformers import T5TokenizerFast
+from transformers import AutoTokenizer
+import os
+
+def load_tokenizer(tokenizer_name, model, dir_path, file_name, model_name = None):
+    
+    if model == "nllb":
+        
+        if not model_name is None:
+        
+            tokenizer = AutoTokenizer.from_pretrained(model_name)
+            
+            print(f"The {model}'s tokenizer was successfully loaded")
+        
+        else:
+            
+            raise ValueError("For the nllb model you must specify the path to the model!")
+    
+    if tokenizer_name == "bpe":
+        
+        tokenizer_path = os.path.join(dir_path, f"{file_name}.json")
+        
+        if model in ["bart", "lstm"]:
+        
+            tokenizer = BartTokenizerFast(tokenizer_file=tokenizer_path)
+        
+            print(f"The Byte Pair Encoding tokenizer was successfully uploaded from {tokenizer_path}")
+        
+    elif tokenizer_name == "sp":
+        
+        tokenizer_path = os.path.join(dir_path, f"{file_name}.model")
+        
+        if model in ['t5', 'mt5']:
+            
+            tokenizer = T5TokenizerFast(vocab_file=tokenizer_path)
+    
+            print(f"The Sentence Piece tokenizer was successfully uploaded from {tokenizer_path}")
+    
+    return tokenizer

translate_package/tokenization/train_tokenizer.py

@@ -0,0 +1,83 @@
+
+from translate_package import argparse, spm, pd, os, Tokenizer, models, pre_tokenizers, trainers, processors, decoders
+from translate_package.errors import TokenizerException
+
+def train_tokenizer(arguments):
+    
+    # recuperate dataset
+    dataset = pd.read_csv(arguments.dataset_file).astype(str)
+    
+    if arguments.name == 'bpe':
+        
+        # initialize tokenizer
+        tokenizer = Tokenizer(models.BPE())
+        
+        # initialize pre-to
+        tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space = False)
+        
+        # initialize trainer
+        trainer = trainers.BpeTrainer(vocab_size = arguments.vocab_size, special_tokens = ['<s>', '<pad>', '</s>', '<unk>'])
+        
+        # iterate over the dataset and return the sentences
+        def get_training_corpus():
+            
+            sentences = dataset[arguments.src_label].tolist() + dataset[arguments.tgt_label].tolist()
+            
+            for i in range(0, len(sentences), arguments.batch_size):
+                
+                yield sentences[i: i + arguments.batch_size]
+                
+        # train the tokenizer
+        tokenizer.train_from_iterator(get_training_corpus(), trainer = trainer)
+        
+        tokenizer.post_processor = processors.ByteLevel(trim_offsets = False)
+        
+        tokenizer.decoder = decoders.ByteLevel()
+        
+        # get path
+        tk_path = os.path.join(arguments.save_path, f'{arguments.file_name}.json')
+        
+        tokenizer.save(tk_path)
+        
+        print(f"The Byte Pair Encoding tokenizer was saved as {tk_path}!")
+    
+    elif arguments.name == 'sp':
+        
+        # print sentences into a file
+        with open('sents.txt', 'w', encoding = 'utf-8') as f:
+            
+            sentences =  dataset[arguments.src_label].tolist() + dataset[arguments.tgt_label].tolist()
+            
+            for i in range(0, len(sentences), arguments.batch_size):
+                
+                sents = sentences[i: i + arguments.batch_size]
+                
+                f.write("\n".join(sents)+'\n')
+        
+        # get path
+        tk_path = os.path.join(arguments.save_path, arguments.file_name)
+              
+        # initialize sentence piece trainer
+        spm.SentencePieceTrainer.Train(input = f'sents.txt',
+                               model_prefix=os.path.join(arguments.save_path, arguments.file_name),
+                               vocab_size=arguments.vocab_size,
+                               character_coverage=1.0,
+                               pad_id=0,                
+                               eos_id=1,
+                               unk_id=2,
+                               bos_id=3,
+                               pad_piece='<pad>',
+                               eos_piece='</s>',
+                               unk_piece='<unk>',
+                               bos_piece='<s>',
+                              )
+        
+        # remove file
+        os.remove('sents.txt')
+        
+        print(f"The Sentence Piece tokenizer was saved as {tk_path}(.model / for model)!")
+        
+    else:
+        
+        raise TokenizerException("You can only train a sentence piece (as 'sp') tokenizer, or a byte pair encoding (as 'bpe') tokenizer!")
+    

translate_package/utils/checkpoint.py

@@ -0,0 +1,16 @@
+from translate_package import wandb
+
+
+def download_checkpoint(project, artifact_path, key):
+
+    wandb.login(key=key)
+
+    run = wandb.init(project=project)
+
+    artifact = run.use_artifact(artifact_path, type="dataset")
+
+    artifact_dir = artifact.download()
+
+    wandb.finish()
+
+    return artifact_dir

translate_package-0.0.1.dist-info/METADATA

@@ -0,0 +1,31 @@
+Metadata-Version: 2.1
+Name: translate-package
+Version: 0.0.1
+Summary: Contain functions and classes to efficiently train a sequence to sequence to translate between two languages.
+Author: Oumar Kane
+Author-email: oumar.kane@univ-thies.sn
+Requires-Dist: accelerate
+Requires-Dist: torch
+Requires-Dist: spacy
+Requires-Dist: nltk
+Requires-Dist: gensim
+Requires-Dist: furo
+Requires-Dist: streamlit
+Requires-Dist: tokenizers
+Requires-Dist: tensorboard
+Requires-Dist: evaluate
+Requires-Dist: transformers
+Requires-Dist: pandas
+Requires-Dist: numpy
+Requires-Dist: scikit-learn
+Requires-Dist: matplotlib
+Requires-Dist: plotly
+Requires-Dist: sacrebleu
+Requires-Dist: nlpaug
+Requires-Dist: wandb
+Requires-Dist: pytorch-lightning
+Requires-Dist: selenium
+Requires-Dist: sentencepiece
+Requires-Dist: peft
+Requires-Dist: rouge-score
+

translate_package-0.0.1.dist-info/RECORD

@@ -0,0 +1,17 @@
+translate_package/__init__.py,sha256=Nckjm15LBEfKSU5-EBjfRkHJQhWELzwow0BD3rKtmkw,1297
+translate_package/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+translate_package/data/data_preparation.py,sha256=sDMKL9LWcwXp6Iy2eS-1n-NWgkZhvhRyFN7k7uID1_w,14804
+translate_package/errors/__init__.py,sha256=gu6XjAIghG4lLkYo8x_7_yyLRtK2FIvmC-WcfJaeOlg,299
+translate_package/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+translate_package/models/gradient_observation.py,sha256=P91UA5i-RdkK46TqpPOJ54DsUYgTI9cRohgPS1Ch0Lc,294
+translate_package/models/lstm.py,sha256=OPkvvceowz5JqdGGH4cfPhH23kbP11z-29zIJn5d8ig,3273
+translate_package/models/machine_translation.py,sha256=5QQpjs_HR9mnPryMyfYpcMgU5tHAAj-eVrv3oGmjR5Y,9963
+translate_package/tokenization/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+translate_package/tokenization/load_tokenizer.py,sha256=vzCHS0ZDSJyr0y08zNvupMtD2jP8A16EBN-ob0LJHG0,1344
+translate_package/tokenization/train_tokenizer.py,sha256=RkdT5DUx201OBNaswM6m54iqcrmCThd3ITLguQb_zVM,3347
+translate_package/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+translate_package/utils/checkpoint.py,sha256=GqymRvF8_QZgrQq9m79Ppj6Qr7NQm78kDARm3p_chC0,322
+translate_package-0.0.1.dist-info/METADATA,sha256=eHD4CGRhoQHOl-EXUPYPTQ9kGQCw5q3Ju0xkmB3PGTM,819
+translate_package-0.0.1.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
+translate_package-0.0.1.dist-info/top_level.txt,sha256=8e2HIrGAMzoSukqu2q929dOJMV1zGYKI_BAFwl-P7XU,18
+translate_package-0.0.1.dist-info/RECORD,,

translate_package-0.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: bdist_wheel (0.37.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

translate_package-0.0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+translate_package