Upload src/ with huggingface_hub

src/MLM/__init__.py

@@ -0,0 +1,2 @@
+from .training_scripts.train_with_trainer import train_with_trainer
+from .datasets.preprocess_dataset import preprocess_dataset

src/MLM/datasets/MLMDataset.py

@@ -0,0 +1,193 @@
+import os
+from dotenv import load_dotenv
+import pandas as pd
+from sklearn.model_selection import train_test_split
+from transformers import BertTokenizerFast, AutoTokenizer
+from datasets import Dataset, DatasetDict, load_dataset
+
+from src.utils import (
+    detect_language,
+    add_emoji_tokens,
+    add_new_line_token,
+    user_id,
+)
+from src.utils.text_functions import clean_text
+from src.utils.s3 import read_csv, save_csv
+
+load_dotenv()
+
+
+class MLMDataset:
+    def __init__(
+        self,
+        s3: bool = False,
+        bucket: str = "lebesgue-data-science",
+        folder: str = os.getenv("GLOBAL_PATH_TO_REPO") + "/data/pretrain",
+        s3_folder: str = "transformers/data/pretrain",
+    ):
+        self.s3 = s3
+        self.bucket = bucket
+
+        if self.s3:
+            self.folder = s3_folder
+        else:
+            self.folder = folder
+
+        self.primaries_path = f"{self.folder}/primaries.csv"
+        self.competitors_path = f"{self.folder}/competitor_ads.csv"
+        self.ad_copies_path = f"{self.folder}/ad_copies.csv"
+        self.english_copies_path = f"{self.folder}/english_copies.csv"
+        self.train_path = f"{self.folder}/train.csv"
+        self.val_path = f"{self.folder}/val.csv"
+        self.test_path = f"{self.folder}/test.csv"
+
+        self.tokenizer_id = f"{user_id}/lebesgue_ad_tokenizer"
+
+        self.hub_datasetdict_id = f"{user_id}/lebesgue_ad_datasets"
+
+    @property
+    def primaries(self) -> pd.DataFrame:
+        df = read_csv(self.primaries_path, s3=self.s3, s3_args={"bucket": self.bucket})
+        return df
+
+    @property
+    def competitors(self) -> pd.DataFrame:
+        df = read_csv(self.competitors_path, s3=self.s3, s3_args={"bucket": self.bucket})
+        return df
+
+    @property
+    def ad_copies(self) -> pd.DataFrame:
+        df = read_csv(self.ad_copies_path, s3=self.s3, s3_args={"bucket": self.bucket})
+        return df
+
+    @property
+    def english_copies(self) -> pd.DataFrame:
+        args = {"lineterminator": "\n"}
+        df = read_csv(
+            self.english_copies_path,
+            s3=self.s3,
+            s3_args={"bucket": self.bucket} | args,
+            pd_args=args,
+        )
+        return df
+
+    @property
+    def train(self) -> pd.DataFrame:
+        df = read_csv(self.train_path, s3=self.s3, s3_args={"bucket": self.bucket})
+        return df
+
+    @property
+    def val(self) -> pd.DataFrame:
+        df = read_csv(self.val_path, s3=self.s3, s3_args={"bucket": self.bucket})
+        return df
+
+    @property
+    def test(self) -> pd.DataFrame:
+        df = read_csv(self.test_path, s3=self.s3, s3_args={"bucket": self.bucket})
+        return df
+
+    @property
+    def datasets(self) -> DatasetDict:
+        return load_dataset(self.hub_datasetdict_id)
+
+    def tokenizer(self, checkpoint: str = "bert-base-uncased") -> AutoTokenizer:
+
+        return AutoTokenizer.from_pretrained(f"{self.tokenizer_id}_{checkpoint}")
+
+    def concat_and_remove_duplicates(self) -> pd.DataFrame:
+
+        comp = self.competitors
+        prim = self.primaries
+
+        primaries = prim.value.to_list()
+        primaries = [primary for primary in primaries if type(primary) == list]
+
+        list_of_primaries = []
+        for primary in primaries:
+            list_of_primaries.extend(primary)
+
+        competitors = comp.ad_text.to_list()
+
+        ad_copies = list_of_primaries + competitors
+        ad_copies = pd.Series(ad_copies).drop_duplicates()
+        ad_copies = pd.DataFrame(ad_copies, columns=["text"])
+        save_csv(
+            df=ad_copies,
+            path=self.ad_copies_path,
+            s3=self.s3,
+            s3_args={"bucket": self.bucket},
+        )
+
+    def get_language(self) -> pd.DataFrame:
+        ad_copies = self.ad_copies
+        ad_copies["language"] = ad_copies.text.apply(lambda text: detect_language(text))
+        save_csv(
+            df=ad_copies,
+            path=self.ad_copies_path,
+            s3=self.s3,
+            s3_args={"bucket": self.bucket},
+        )
+        return ad_copies
+
+    def filter_english(self) -> pd.DataFrame:
+        ad_copies = self.ad_copies
+        english = ad_copies[ad_copies.language == "en"]
+        save_csv(
+            df=english,
+            path=self.english_copies_path,
+            s3=self.s3,
+            s3_args={"bucket": self.bucket},
+        )
+        return english
+
+    def clean_english(self) -> pd.DataFrame:
+        english = self.english_copies
+        english["text_clean"] = english.text.apply(clean_text)
+
+        # remove empty ones
+        english = english[english.text_clean.apply(len) != 0]
+        save_csv(
+            df=english,
+            path=self.english_copies_path,
+            s3=self.s3,
+            s3_args={"bucket": self.bucket},
+        )
+        return english
+
+    def train_tokenizer(self, checkpoint: str = "bert-base-uncased"):
+
+        tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased")
+        tokenizer = add_emoji_tokens(tokenizer=tokenizer)
+        tokenizer = add_new_line_token(tokenizer=tokenizer)
+
+        tokenizer.push_to_hub(f"{self.tokenizer_id}_{checkpoint}")
+
+    def get_tokenizer(self):
+        return BertTokenizerFast.from_pretrained(self.tokenizer_id)
+
+    def split_into_train_and_test(
+        self,
+    ) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
+        df = self.english_copies
+        train, test = train_test_split(df, train_size=0.9, random_state=42)
+        train, val = train_test_split(train, train_size=0.85, random_state=42)
+
+        dataset_dict = DatasetDict()
+
+        for df, local_path, dataset_dict_key in zip(
+            [train, val, test],
+            [self.train_path, self.val_path, self.train_path],
+            ["train", "val", "test"],
+        ):
+            save_csv(df=df, path=local_path, s3=self.s3, s3_args={"bucket": self.bucket})
+            df_hf = Dataset.from_pandas(df, preserve_index=False)
+            dataset_dict[dataset_dict_key] = df_hf
+
+        dataset_dict.push_to_hub(self.hub_datasetdict_id)
+
+        return train, val, test
+
+
+mlm_dataset = MLMDataset()
+
+mlm_dataset_s3 = MLMDataset(s3=True)

src/MLM/datasets/__init__.py

@@ -0,0 +1,2 @@
+from .preprocess_dataset import preprocess_dataset
+from .MLMDataset import MLMDataset, mlm_dataset, mlm_dataset_s3

src/MLM/datasets/preprocess_dataset.py

@@ -0,0 +1,39 @@
+from datasets import Dataset, DatasetDict
+from transformers import AutoTokenizer
+
+def preprocess_dataset(dataset: Dataset | DatasetDict, tokenizer: AutoTokenizer) -> Dataset | DatasetDict:
+
+    tokenized_dataset = dataset.map(
+        lambda examples: tokenize_function(examples, tokenizer), batched=True, remove_columns=["text", 'text_clean', 'language'] 
+    )
+
+    return tokenized_dataset.map(group_texts, batched=True)
+
+
+
+
+
+
+
+def tokenize_function(examples, tokenizer: AutoTokenizer):
+
+    result = tokenizer(examples["text"])
+
+    if tokenizer.is_fast:
+        result["word_ids"] = [result.word_ids(i) for i in range(len(result["input_ids"]))]
+    
+    return result
+
+
+def group_texts(examples, chunk_size: int = 128):
+
+    concatinated_examples = {k : sum(examples[k], []) for k in examples.keys()}
+
+    total_length = len(concatinated_examples["input_ids"])
+    total_length = (total_length // chunk_size) * chunk_size
+
+    result = {k : [t[i : i+chunk_size] for i in range(0, total_length, chunk_size)] for k, t in concatinated_examples.items()}
+    result["labels"] = result["input_ids"].copy()
+
+    return result
+

src/MLM/mask_and_unmask.py

@@ -0,0 +1,32 @@
+from transformers import AutoTokenizer, AutoModelForMaskedLM, BertTokenizerFast, DataCollatorForLanguageModeling
+import torch
+
+
+def mask_and_unmask(
+    text: str,
+    tokenizer: AutoTokenizer | BertTokenizerFast,
+    model: AutoModelForMaskedLM,
+    data_collator: DataCollatorForLanguageModeling,
+) -> str:
+
+    collator_input = tokenizer(text)
+    collator_input["labels"] = collator_input["input_ids"].copy()
+    collator_output = data_collator([collator_input])
+    masked_text = tokenizer.decode(collator_output["input_ids"][0])
+
+    pred_dict = {"masked_text": masked_text}
+
+    inputs = tokenizer(masked_text, return_tensors="pt", padding="max_length", truncation=True)
+    token_logits = model(input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"]).logits
+    all_masked_token_index = torch.argwhere(inputs["input_ids"] == tokenizer.mask_token_id)
+    if all_masked_token_index.size()[0] != 0:
+
+        for i, masked_index_token in enumerate(all_masked_token_index[:, 1]):
+            # print(masked_index_token)
+            masked_token_logits = token_logits[0, masked_index_token, :]
+            # print(masked_token_logits)
+            top_5_tokens = torch.argsort(masked_token_logits, descending=True)[:5].tolist()
+            value = tokenizer.decode(collator_output["labels"][0, masked_index_token - 1])
+            pred_dict[value] = [tokenizer.decode(token) for token in top_5_tokens]
+
+    return pred_dict

src/MLM/training_scripts/__init__.py

@@ -0,0 +1 @@
+from .train_with_trainer import train_with_trainer

src/MLM/training_scripts/train_with_trainer.py

@@ -0,0 +1,50 @@
+from transformers import AutoModelForMaskedLM, AutoTokenizer, TrainingArguments, Trainer
+from datasets import Dataset, DatasetDict
+from transformers import DataCollatorForLanguageModeling
+
+from src.MLM.datasets.preprocess_dataset import preprocess_dataset
+from src.MLM.training_scripts.utils import get_new_model_name
+
+
+def train_with_trainer(
+    model_checkpoint: str,
+    tokenizer: AutoTokenizer,
+    dataset: DatasetDict,
+    model_name: str | None = None,
+    data_collator=None,
+    num_epochs: int = 3,
+):
+
+    model = AutoModelForMaskedLM.from_pretrained(model_checkpoint)
+
+    model_name = get_new_model_name(model_checkpoint=model_checkpoint, model_name=model_name)
+
+    dataset = preprocess_dataset(dataset=dataset, tokenizer=tokenizer)
+
+    if data_collator is None:
+        data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=0.15)
+
+    training_args = TrainingArguments(
+        model_name,
+        evaluation_strategy="epoch",
+        learning_rate=2e-5,
+        weight_decay=0.01,
+        push_to_hub=True,
+        report_to="wandb",
+        run_name=model_name,
+        num_train_epochs=num_epochs,
+        save_total_limit=1,
+        save_strategy="epoch",
+    )
+
+    print(f"device: {training_args.device}")
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=dataset["train"],
+        eval_dataset=dataset["val"],
+        data_collator=data_collator,
+    )
+
+    trainer.train()

src/MLM/training_scripts/utils.py

@@ -0,0 +1,9 @@
+def get_new_model_name(model_checkpoint: str, model_name: str = None) -> str:
+    if model_name is None:
+        old_version_number = int(model_checkpoint[-2:])
+        new_version_number = str(old_version_number + 1).zfill(2)
+        model_name = f"{model_checkpoint[:-2]}{new_version_number}"
+    elif not model_name[-2:].isnumeric():
+        model_name = model_name + "_00"
+
+    return model_name

src/regression/.gitignore

@@ -0,0 +1 @@
+runs/

src/regression/HF/__init__.py

@@ -0,0 +1,2 @@
+from .configs import *
+from .models import *

src/regression/HF/configs/FullModelConfigHF.py

@@ -0,0 +1,25 @@
+from transformers import PretrainedConfig
+from src.regression.PL import EncoderPL, DecoderPL
+from typing import List
+
+
+class FullModelConfigHF(PretrainedConfig):
+
+    model_type = "full_model"
+
+    def __init__(
+        self,
+        tokenizer_ckpt: str = "",
+        bert_ckpt: str = "",
+        decoder_ckpt: str = "",
+        layer_norm: bool = True,
+        nontext_features: List[str] = ["aov"],
+        **kwargs,
+    ):
+
+        self.tokenizer_ckpt = tokenizer_ckpt
+        self.bert_ckpt = bert_ckpt
+        self.decoder_ckpt = decoder_ckpt
+        self.nontext_features = nontext_features
+        self.layer_norm = layer_norm
+        super().__init__(**kwargs)

src/regression/HF/configs/__init__.py

@@ -0,0 +1 @@
+from .FullModelConfigHF import FullModelConfigHF

src/regression/HF/models/FullModelHF.py

@@ -0,0 +1,43 @@
+from transformers import PreTrainedModel
+from transformers import AutoModelForMaskedLM, AutoTokenizer
+from pytorch_lightning.loggers import WandbLogger
+
+from src.regression.PL import FullModelPL, EncoderPL, DecoderPL
+from src.regression.HF.configs import FullModelConfigHF
+
+from config import DEVICE
+
+
+class FullModelHF(PreTrainedModel):
+    config_class = FullModelConfigHF
+
+    def __init__(self, config):
+
+        super().__init__(config)
+
+        self.tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_ckpt)
+        mlm_bert = AutoModelForMaskedLM.from_pretrained(config.bert_ckpt)
+        self.bert = mlm_bert.distilbert
+
+        encoder = EncoderPL(tokenizer=self.tokenizer, bert=self.bert).to(DEVICE)
+
+        wandb_logger = WandbLogger(
+            project="transformers",
+            entity="sanjin_juric_fot",
+            # log_model=True,
+            # reinit=True,
+        )
+
+        artifact = wandb_logger.use_artifact(config.decoder_ckpt)
+        artifact_dir = artifact.download()
+        decoder = DecoderPL.load_from_checkpoint(artifact_dir + "/" + "model.ckpt").to(DEVICE)
+
+        self.model = FullModelPL(
+            encoder=encoder,
+            decoder=decoder,
+            layer_norm=config.layer_norm,
+            nontext_features=config.nontext_features,
+        ).to(DEVICE)
+
+    def forward(self, input):
+        return self.model._get_loss(input)

src/regression/HF/models/__init__.py

@@ -0,0 +1 @@
+from .FullModelHF import FullModelHF

src/regression/PL/DecoderPL.py

@@ -0,0 +1,180 @@
+import emoji
+import numpy as np
+import pytorch_lightning as pl
+import torch
+import torch.nn.functional as F
+from loguru import logger
+from torch import nn
+from torch.optim.lr_scheduler import CosineAnnealingLR
+from torchmetrics import R2Score
+
+from src.utils import get_sentiment
+from src.utils.neural_networks import set_layer
+from config import DEVICE
+
+torch.set_default_dtype(torch.float32)
+
+
+class DecoderPL(pl.LightningModule):
+    def __init__(
+        self,
+        input_dim: int = 774,
+        layer_norm: bool = True,
+        layer_dict: dict = {},
+        device=DEVICE,
+        T_max: int = 10,
+        start_lr: float = 5 * 1e-4,
+    ):
+        super().__init__()
+
+        # layers
+        self.linear1 = set_layer(
+            layer_dict=layer_dict,
+            name="linear1",
+            alternative=nn.Linear(in_features=input_dim, out_features=512),
+        )
+
+        self.linear2 = set_layer(
+            layer_dict=layer_dict,
+            name="linear2",
+            alternative=nn.Linear(in_features=512, out_features=264),
+        )
+
+        self.linear3 = set_layer(
+            layer_dict=layer_dict,
+            name="linear3",
+            alternative=nn.Linear(in_features=264, out_features=64),
+        )
+
+        self.linear4 = set_layer(
+            layer_dict=layer_dict,
+            name="linear4",
+            alternative=nn.Linear(in_features=64, out_features=1),
+        )
+
+        self.activation = nn.LeakyReLU(negative_slope=0.01)
+
+        if not layer_norm:
+            self.layers = [
+                self.linear1,
+                self.activation,
+                self.linear2,
+                self.activation,
+                self.linear3,
+                self.activation,
+                self.linear4,
+            ]
+        else:
+            self.layernorm1 = nn.LayerNorm(normalized_shape=(1, self.linear1.out_features))
+            self.layernorm2 = nn.LayerNorm(normalized_shape=(1, self.linear2.out_features))
+            self.layernorm3 = nn.LayerNorm(normalized_shape=(1, self.linear3.out_features))
+            self.layers = [
+                self.linear1,
+                self.layernorm1,
+                self.activation,
+                self.linear2,
+                self.layernorm2,
+                self.activation,
+                self.linear3,
+                self.layernorm3,
+                self.activation,
+                self.linear4,
+            ]
+
+        # initialize weights
+        [self.initialize_weights(layer) for layer in self.layers]
+
+        # optimizer and scheduler
+        self.optimizer = torch.optim.AdamW(filter(lambda p: p.requires_grad, self.parameters()), lr=start_lr)
+        self.scheduler = CosineAnnealingLR(self.optimizer, T_max=T_max)
+
+        # else
+        self.save_hyperparameters(ignore=["model"])
+        self.MSE = nn.MSELoss()
+        self.R2 = R2Score()
+
+    def initialize_weights(self, module):
+
+        if isinstance(module, nn.Linear):
+            logger.debug("linear weights initialized")
+            torch.nn.init.xavier_uniform_(module.weight)
+            module.bias.data.fill_(0.01)
+
+    def forward(self, x: torch.Tensor):
+
+        if x.dim() == 2:
+            x = x.unsqueeze(dim=1)
+
+        for layer in self.layers:
+            x = layer(x)
+
+        x = x.squeeze()
+
+        if x.dim() == 0:
+            x = x.unsqueeze(dim=0)
+
+        return x.to(torch.float32)
+
+    def training_step(self, batch):
+
+        loss_and_metrics = self._get_loss(batch, get_metrics=True)
+        pred = loss_and_metrics["pred"]
+        act = loss_and_metrics["act"]
+        loss = loss_and_metrics["loss"]
+
+        self.log("train_loss", loss, on_epoch=True, on_step=False, prog_bar=True, logger=True)
+
+        return {"loss": loss, "pred": pred, "act": act}
+
+    def configure_optimizers(self):
+
+        optimizer = self.optimizer
+        scheduler = self.scheduler
+        return dict(optimizer=optimizer, lr_scheduler=scheduler)
+
+    def lr_scheduler_step(self, scheduler, optimizer_idx, metric):
+        logger.debug(scheduler)
+        if metric is None:
+            scheduler.step()
+        else:
+            scheduler.step(metric)
+
+    def validation_step(self, batch, batch_idx):
+        """used for logging metrics"""
+        loss_and_metrics = self._get_loss(batch, get_metrics=True)
+        loss = loss_and_metrics["loss"]
+
+        # Log loss and metric
+        self.log("val_loss", loss, on_epoch=True, prog_bar=True, logger=True)
+
+    def training_epoch_end(self, training_step_outputs):
+
+        training_step_outputs = list(training_step_outputs)
+
+        training_step_outputs.pop()
+
+        output_dict = {k: [dic[k] for dic in training_step_outputs] for k in training_step_outputs[0]}
+
+        pred = torch.stack(output_dict["pred"])
+        act = torch.stack(output_dict["act"])
+
+        loss = torch.sub(pred, act)
+        loss_sq = torch.square(loss)
+
+        TSS = float(torch.var(act, unbiased=False))
+        RSS = float(torch.mean(loss_sq))
+        R2 = 1 - RSS / TSS
+
+        self.log("train_R2", R2, prog_bar=True, logger=True)
+
+    def _get_loss(self, batch, get_metrics: bool = False):
+        """convenience function since train/valid/test steps are similar"""
+        pred = self.forward(x=batch["embedding"]).to(torch.float32)
+
+        act, loss = None, None
+
+        if "ctr" in batch.keys():
+            act = batch["ctr"].to(torch.float32)
+            loss = self.MSE(pred, act).to(torch.float32)
+
+        return {"loss": loss, "pred": pred, "act": act}

src/regression/PL/EncoderPL.py

@@ -0,0 +1,116 @@
+import emoji
+import numpy as np
+import pytorch_lightning as pl
+import torch
+import torch.nn.functional as F
+from loguru import logger
+from torch import nn
+from torch.optim.lr_scheduler import CosineAnnealingLR
+from torchmetrics import R2Score
+from transformers import BertModel, BertTokenizer, DistilBertModel, AutoModel, AutoTokenizer
+from pytorch_lightning import LightningModule
+
+
+from src.utils.neural_networks import set_layer
+from src.utils import add_emoji_tokens, add_new_line_token, vectorise_dict
+from config import DEVICE
+
+torch.set_default_dtype(torch.float32)
+
+
+class EncoderPL(pl.LightningModule):
+    def __init__(
+        self,
+        model_name: str = "bert-base-uncased",
+        tokenizer: AutoTokenizer | None = None,
+        bert: AutoModel | None = None,
+        cls: bool = False,
+        device=DEVICE,
+    ):
+        super().__init__()
+
+        self._device = device
+        self.cls = cls
+        self.model_name = model_name
+
+        # layers
+
+        self.tokenizer = tokenizer if tokenizer is not None else BertTokenizer.from_pretrained(model_name)
+
+        self.bert = bert if bert is not None else BertModel.from_pretrained(model_name)
+
+        if tokenizer is None:
+            self.tokenizer = add_emoji_tokens(self.tokenizer)
+            self.tokenizer = add_new_line_token(self.tokenizer)
+            self.bert.resize_token_embeddings(len(self.tokenizer))
+
+        # optimizer and scheduler
+        self.optimizer = torch.optim.AdamW(filter(lambda p: p.requires_grad, self.parameters()), lr=1e-3)
+
+        # config tweaking
+        self.bert.config.torch_dtype = "float32"
+
+    def forward(self, text: str):
+
+        # run text through bert and squash the output to get embeddings
+        encoded = self.tokenizer(text, return_tensors="pt", padding="max_length", truncation=True).to(self._device)
+
+        if type(self.bert) == DistilBertModel:
+            encoded.pop("token_type_ids")
+
+        bert_output = self.bert(**encoded)
+
+        if self.cls:
+            if hasattr(bert_output, "pooler_output") and bert_output.pooler_output is not None:
+                embedding = bert_output.pooler_output.unsqueeze(dim=1)
+            else:
+                embedding = bert_output.last_hidden_state[0, 0, :].unsqueeze(dim=0).unsqueeze(dim=0)
+        else:
+            last_hidden_state = bert_output.last_hidden_state
+
+            if last_hidden_state.dim() == 2:
+                last_hidden_state = last_hidden_state.unsqueeze(dim=0)
+
+            embedding = torch.matmul(
+                encoded["attention_mask"].type(torch.float32).view(-1, 1, 512),
+                last_hidden_state,
+            )
+
+        return embedding
+
+    def configure_optimizers(self):
+        return self.optimizer
+
+
+def get_bert_embedding(
+    text: str, as_list: bool = True, cls: bool = False, device=DEVICE, layer_dict: dict = {}
+) -> list:
+    encoder = EncoderPL(cls=cls, layer_dict=layer_dict).to(device)
+    embedding = encoder.forward(text)
+
+    if as_list:
+        embedding = embedding.tolist()[0][0]
+
+    return embedding
+
+
+def get_concat_embedding(
+    text: str = None,
+    bert_embedding: list = [],
+    other_features: dict = {},
+    cls: bool = False,
+    device=DEVICE,
+    layer_dict: dict = {},
+) -> list:
+
+    if not len(bert_embedding):
+
+        if text is None:
+            raise ValueError("both text and embedding are empty!")
+        bert_embedding = get_bert_embedding(text=text, cls=cls, device=device, layer_dict=layer_dict)
+
+    other_features = vectorise_dict(other_features, as_list=True)
+
+    concat_vec = bert_embedding + other_features
+
+    return concat_vec

src/regression/PL/FullModelPL.py

@@ -0,0 +1,166 @@
+import emoji
+import numpy as np
+import pytorch_lightning as pl
+import torch
+import torch.nn.functional as F
+from loguru import logger
+from torch import nn
+from torch.optim.lr_scheduler import CosineAnnealingLR
+from torchmetrics import R2Score
+from transformers import BertModel, BertTokenizerFast
+
+from src.utils import get_sentiment, vectorise_dict
+from src.utils.neural_networks import set_layer
+from config import DEVICE
+
+from .DecoderPL import DecoderPL
+from .EncoderPL import EncoderPL
+
+
+torch.set_default_dtype(torch.float32)
+
+
+class FullModelPL(pl.LightningModule):
+    def __init__(
+        self,
+        model_name: str = "bert-base-uncased",
+        nontext_features: list[str] = ["aov"],
+        encoder: EncoderPL | None = None,
+        decoder: DecoderPL | None = None,
+        layer_norm: bool = True,
+        device=DEVICE,
+        T_max: int = 10,
+    ):
+        super().__init__()
+
+        # layers
+        self.encoder = (
+            encoder.to(self.device)
+            if encoder is not None
+            else EncoderPL(model_name=model_name, device=device).to(self.device)
+        )
+        self.decoder = (
+            decoder.to(self.device)
+            if decoder is not None
+            else DecoderPL(
+                input_dim=768 + len(nontext_features) + 5,
+                layer_norm=layer_norm,
+                device=device,
+            ).to(self.device)
+        )
+
+        # else
+        self.MSE = nn.MSELoss()
+        self.R2 = R2Score()
+
+        self.optimizer = torch.optim.AdamW(filter(lambda p: p.requires_grad, self.parameters()), lr=3 * 1e-4)
+        self.scheduler = CosineAnnealingLR(self.optimizer, T_max=T_max)
+
+        # self.save_hyperparameters(ignore=["decoder", "encoder"])
+
+    def forward(self, input_dict: dict):
+
+        input_dict = input_dict.copy()
+        text = input_dict.pop("text")
+
+        print(f"text: {text}")
+
+        if "ctr" in input_dict.keys():
+            input_dict.pop("ctr")
+
+        # encode
+        sentence_embedding = self.encoder.forward(text=text)
+
+        # sentiment
+        sentiment = get_sentiment_for_list_of_texts(text)
+        input_dict = input_dict | sentiment
+
+        input_dict = {k: v.to(self.device) for k, v in input_dict.items()}
+
+        # concat nontext features to embedding
+        nontext_vec = vectorise_dict(input_dict)
+        nontext_tensor = torch.stack(nontext_vec).T.unsqueeze(1).to(torch.float32)
+        # logger.debug(f"nontext tensor type: {nontext_tensor.dtype}")
+        print(f"{sentence_embedding.get_device()}, {nontext_tensor.get_device()}")
+        x = torch.cat((sentence_embedding, nontext_tensor), 2)
+
+        print(self.decoder.device)
+        print(x.get_device())
+
+        # decode
+        result = self.decoder.forward(x)
+        return result
+
+    def training_step(self, batch):
+
+        loss_and_metrics = self._get_loss(batch, get_metrics=True)
+        pred = loss_and_metrics["pred"]
+        act = loss_and_metrics["act"]
+        loss = loss_and_metrics["loss"]
+
+        self.log("train_loss", loss, on_epoch=True, on_step=False, prog_bar=True, logger=True)
+
+        return {"loss": loss, "pred": pred, "act": act}
+
+    def configure_optimizers(self):
+
+        for name, param in self.named_parameters():
+            if "bert" in name:
+                param.requires_grad = False
+
+        optimizer = self.optimizer
+        scheduler = self.scheduler
+        return dict(optimizer=optimizer, lr_scheduler=scheduler)
+
+    def lr_scheduler_step(self, scheduler, optimizer_idx, metric):
+        logger.debug(scheduler)
+        if metric is None:
+            scheduler.step()
+        else:
+            scheduler.step(metric)
+
+    def validation_step(self, batch, batch_idx):
+        """used for logging metrics"""
+        loss_and_metrics = self._get_loss(batch, get_metrics=True)
+        loss = loss_and_metrics["loss"]
+
+        # Log loss and metric
+        self.log("val_loss", loss, on_epoch=True, prog_bar=True, logger=True)
+
+    def training_epoch_end(self, training_step_outputs):
+
+        training_step_outputs = list(training_step_outputs)
+
+        training_step_outputs.pop()
+
+        output_dict = {k: [dic[k] for dic in training_step_outputs] for k in training_step_outputs[0]}
+
+        pred = torch.stack(output_dict["pred"])
+        act = torch.stack(output_dict["act"])
+
+        loss = torch.sub(pred, act)
+        loss_sq = torch.square(loss)
+
+        TSS = float(torch.var(act, unbiased=False))
+        RSS = float(torch.mean(loss_sq))
+        R2 = 1 - RSS / TSS
+
+        self.log("train_R2", R2, prog_bar=True, logger=True)
+
+    def _get_loss(self, batch, get_metrics: bool = False):
+        """convenience function since train/valid/test steps are similar"""
+        pred = self.forward(input_dict=batch).to(torch.float32)
+
+        act, loss = None, None
+
+        if "ctr" in batch.keys():
+            act = batch["ctr"].to(torch.float32).to(self.device)
+            loss = self.MSE(pred, act).to(torch.float32)
+
+        return {"loss": loss, "pred": pred, "act": act}
+
+
+def get_sentiment_for_list_of_texts(texts: list[str]) -> dict:
+    ld = [get_sentiment(text) for text in texts]
+    v = {k: torch.Tensor([dic[k] for dic in ld]) for k in ld[0]}
+    return v

src/regression/PL/__init__.py

@@ -0,0 +1,3 @@
+from .FullModelPL import FullModelPL
+from .DecoderPL import DecoderPL
+from .EncoderPL import EncoderPL, get_concat_embedding, get_bert_embedding

src/regression/__init__.py

@@ -0,0 +1,3 @@
+from .datasets import *
+from .training_scripts import *
+from .PL import *

src/regression/datasets/DecoderDatasetTorch.py

@@ -0,0 +1,38 @@
+import numpy as np
+import pandas as pd
+import torch
+from torch.utils.data import Dataset
+
+
+class DecoderDatasetTorch(Dataset):
+    """Train dataset."""
+
+    def __init__(self, df: pd.DataFrame, embedding_column: str = "my_full_mean_embedding"):
+        """
+
+        Args:
+            df (pd.DataFrame): dataframe with ads
+            embedding_column (str, optional): Column whose values to output in __get_item__. Defaults to 'full_mean_embedding'.
+        """
+        self.df = df
+        self.embedding_column = embedding_column
+
+        df[[embedding_column, "ctr"]] = df[[embedding_column, "ctr"]].applymap(lambda x: np.float32(x))
+        # df["ctr"] = df["ctr"].astype(np.float32)
+
+    def __len__(self):
+        return len(self.df)
+
+    def __getitem__(self, idx):
+
+        if torch.is_tensor(idx):
+            idx = idx.tolist()
+
+        embedding = self.df.loc[idx, self.embedding_column]
+        ctr = self.df.loc[idx, "ctr"]
+
+        return {"embedding": embedding, "ctr": ctr}
+
+
+# tokenizer = BertTokenizer.from_pretrained("textattack/bert-base-uncased-yelp-polarity")
+# train_dataset = AdDataset(df=dataset.train, tokenizer=tokenizer)

src/regression/datasets/FullModelDatasetTorch.py

@@ -0,0 +1,39 @@
+import numpy as np
+import pandas as pd
+import torch
+from torch.utils.data import DataLoader, Dataset
+
+
+class FullModelDatasetTorch(Dataset):
+    """Train dataset."""
+
+    def __init__(self, df: pd.DataFrame, nontext_features: list[str] = ["aov"]):
+        """
+
+        Args:
+            df (pd.DataFrame): train dataframe
+            nontext_features (list[str]): features to use in training except for text embeddings
+        """
+        self.df = df
+        self.nontext_features = nontext_features
+
+        df[nontext_features + ["ctr"]] = df[nontext_features + ["ctr"]].astype(np.float32)
+
+    def __len__(self):
+        return len(self.df)
+
+    def __getitem__(self, idx):
+
+        if torch.is_tensor(idx):
+            idx = idx.tolist()
+
+        text = self.df.loc[idx, "text_clean"]
+        ctr = self.df.loc[idx, "ctr"]
+
+        nontext_features = {feature: self.df.loc[idx, feature] for feature in self.nontext_features}
+
+        return {"text": text, "ctr": ctr} | nontext_features
+
+
+# tokenizer = BertTokenizer.from_pretrained("textattack/bert-base-uncased-yelp-polarity")
+# train_dataset = AdDataset(df=dataset.train, tokenizer=tokenizer)

src/regression/datasets/RegressionDataset.py

@@ -0,0 +1,313 @@
+import os
+
+import numpy as np
+import pandas as pd
+from dotenv import load_dotenv
+from langdetect import detect
+from loguru import logger
+from sklearn.model_selection import train_test_split
+from time import sleep
+from transformers import BertModel, AutoTokenizer
+from tqdm import tqdm
+import torch
+from config import DEVICE
+
+from src.utils.text_functions import clean_text, detect_language
+from src.utils import (
+    get_sentiment,
+    detect_language,
+)
+
+from src.regression.PL import (
+    get_bert_embedding,
+    get_concat_embedding,
+)
+
+from src.utils.s3 import read_csv, save_csv
+
+
+load_dotenv()
+
+
+class RegressionDataset:
+    def __init__(
+        self,
+        s3: bool = False,
+        bucket: str = "lebesgue-data-science",
+        folder: str = os.getenv("GLOBAL_PATH_TO_REPO") + "/data",
+        s3_folder: str = "transformers/data",
+    ):
+        self.s3 = s3
+        self.bucket = bucket
+
+        if self.s3:
+            self.folder = s3_folder
+        else:
+            self.folder = folder
+
+        self.original_path = f"{self.folder}/original.csv"
+        self.untrimmed_path = f"{self.folder}/untrimmed.csv"
+        self.normalized_path = f"{self.folder}/normalized.csv"
+        self.trimmed_path = f"{self.folder}/trimmed.csv"
+
+        self.train_path = f"{self.folder}/train.csv"
+        self.val_path = f"{self.folder}/val.csv"
+        self.test_path = f"{self.folder}/test.csv"
+
+        self.text_types = ["primary", "title", "description"]
+
+        self.col_func_dict = {
+            "number": len,
+            "len": lambda texts: np.mean([len(text) for text in texts]),
+        }
+
+    @property
+    def original(self) -> pd.DataFrame:
+        df = read_csv(path=self.original_path, s3=self.s3, s3_args={"bucket": self.bucket})
+        return df
+
+    @property
+    def untrimmed(self) -> pd.DataFrame:
+        df = read_csv(path=self.untrimmed_path, s3=self.s3, s3_args={"bucket": self.bucket})
+        return df
+
+    @property
+    def normalized(self) -> pd.DataFrame:
+        df = read_csv(path=self.normalized_path, s3=self.s3, s3_args={"bucket": self.bucket})
+        return df
+
+    @property
+    def trimmed(self) -> pd.DataFrame:
+        df = read_csv(path=self.trimmed_path, s3=self.s3, s3_args={"bucket": self.bucket})
+        return df
+
+    @property
+    def train(self) -> pd.DataFrame:
+        df = read_csv(path=self.train_path, s3=self.s3, s3_args={"bucket": self.bucket})
+        return df
+
+    @property
+    def val(self) -> pd.DataFrame:
+        df = read_csv(path=self.val_path, s3=self.s3, s3_args={"bucket": self.bucket})
+        return df
+
+    @property
+    def test(self) -> pd.DataFrame:
+        df = read_csv(path=self.test_path, s3=self.s3, s3_args={"bucket": self.bucket})
+        return df
+
+    def normalize_untrimmed(self, group_cols: list[str] = ["text", "target", "shop_id"]) -> pd.DataFrame:
+        df = self.untrimmed
+        grouped = df.groupby(group_cols)
+
+        filters_df = grouped.agg({"impr": "sum", "spend": "sum"}).reset_index()
+        ctr = grouped.apply(lambda df: df.link_clicks.sum() / df.impr.sum())
+        ctr_df = pd.DataFrame(ctr, columns=["ctr"]).reset_index()
+        normalised = filters_df.merge(ctr_df, on=group_cols)
+
+        merged = df.merge(normalised, on=group_cols, validate="m:1", suffixes=["___", None])
+        merged.drop(list([col for col in merged.columns if "___" in col]), inplace=True, axis=1)
+        final = merged.drop_duplicates(group_cols)
+        save_csv(
+            df=final,
+            path=self.normalized_path,
+            s3=self.s3,
+            s3_args={"bucket": self.bucket},
+        )
+        return df
+
+    def expand_untrimmed(self, update_existing_columns: bool = False) -> pd.DataFrame:
+
+        df = self.untrimmed
+
+        # normalise target by adset
+        # df["ctr_norm"] = (
+        #     df.groupby(["shop_id", "adset_id"])
+        #     .ctr.transform(lambda x: (x - x.mean()) / x.std())
+        #     .count()
+        # )
+
+        new_col_func_dict = self.col_func_dict
+
+        if not update_existing_columns:
+            new_col_func_dict = {
+                col: fun for col, fun in new_col_func_dict.items() if "primary_" + col not in df.columns
+            }
+
+        # get extra columns
+        for col, func in new_col_func_dict.items():
+            logger.debug(col)
+            for text_type in self.text_types:
+                df[f"{text_type}_{col}"] = df[text_type].apply(func)
+
+        df["has_text"] = df.apply(
+            lambda df: bool(df.primary_number + df.title_number + df.description_number),
+            axis=1,
+        )
+
+        # text columns
+        df = df.apply(_get_text, axis=1)
+        df = df.apply(_get_concatinated_text, axis=1)
+
+        df["language"] = df.text.apply(detect_language)
+        df = df[df.language == "en"]
+        df = df[df.ctr.notna()]
+
+        save_csv(df=df, path=self.untrimmed_path, s3=self.s3, s3_args={"bucket": self.bucket})
+
+        return df
+
+    def trim(self, min_impr: int = 900, min_spend: float = 90) -> pd.DataFrame:
+        df = self.normalized
+        df = df[(df.impr >= min_impr) & (df.spend >= min_spend)]
+        df = df[df.target == "acquisition"]
+        df = df[df.aov.notna()]
+
+        df = df[df.has_text == True]
+
+        save_csv(df=df, path=self.trimmed_path, s3=self.s3, s3_args={"bucket": self.bucket})
+
+        return df
+
+    def expand_trimmed(
+        self, bert: BertModel = None, tokenizer: AutoTokenizer = None, add_bert_embeddings_bool: bool = False
+    ) -> pd.DataFrame:
+        df = self.trimmed
+
+        # clean text
+        for col in ["text", "concat_text"]:
+            df[f"{col}_clean"] = df[col].apply(clean_text)
+
+        df["text_clean_sentiment"] = df.text_clean.apply(get_sentiment)
+
+        if add_bert_embeddings_bool:
+            if tokenizer is None or bert is None:
+                raise ValueError("tokenizer or bert is None")
+            layer_dict = {"bert": bert, "tokenizer": tokenizer}
+            df = add_bert_embeddings(df=df, save_path=self.trimmed_path, layer_dict=layer_dict)
+
+            df = df.apply(add_concat_embeddings, axis=1)
+
+        save_csv(df=df, path=self.trimmed_path, s3=self.s3, s3_args={"bucket": self.bucket})
+        return df
+
+    def split_into_train_and_test(
+        self,
+    ) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
+        df = self.trimmed
+        train, test = train_test_split(df, train_size=0.9, random_state=42)
+        train, val = train_test_split(train, train_size=0.85, random_state=42)
+        save_csv(df=train, path=self.train_path, s3=self.s3, s3_args={"bucket": self.bucket})
+        save_csv(df=val, path=self.val_path, s3=self.s3, s3_args={"bucket": self.bucket})
+        save_csv(df=test, path=self.test_path, s3=self.s3, s3_args={"bucket": self.bucket})
+        return train, val, test
+
+    def expand_normalise_trim_split(
+        self,
+        update_existing_columns: bool = False,
+        group_cols=["text", "target", "shop_id"],
+    ) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
+        self.expand_untrimmed(update_existing_columns=update_existing_columns)
+        self.normalize_untrimmed(group_cols=group_cols)
+        self.trim()
+        self.expand_trimmed()
+        train, val, test = self.split_into_train_and_test()
+        return train, val, test
+
+
+def _get_text(ad: pd.Series) -> pd.Series:
+
+    if ad.primary_number > 0:
+        ad["text"] = ad.primary[0]
+
+    elif ad.description_number > 0:
+        ad["text"] = ad.description[0]
+
+    elif ad.title_number > 0:
+        ad["text"] = ad.title[0]
+
+    else:
+        ad["text"] = None
+
+    return ad
+
+
+def _get_concatinated_text(ad: pd.Series) -> pd.Series:
+
+    concat_text = ""
+
+    if ad.primary_number > 0:
+        concat_text = concat_text + ad.primary[0]
+
+    if ad.description_number > 0:
+        concat_text = concat_text + ad.description[0]
+
+    if ad.title_number > 0:
+        concat_text = concat_text + ad.title[0]
+
+    ad["concat_text"] = concat_text
+
+    return ad
+
+
+regression_dataset = RegressionDataset()
+
+regression_dataset_s3 = RegressionDataset(s3=True)
+
+
+def add_bert_embeddings(df: pd.DataFrame, save_path: str, layer_dict: dict = {}, device=DEVICE) -> pd.DataFrame:
+
+    if device == torch.device("cuda"):
+        df["my_bert_cls_embedding"] = df.text_clean.apply(
+            lambda text: get_bert_embedding(text=text, cls=True, layer_dict=layer_dict)
+        )
+        df["my_bert_mean_embedding"] = df.text_clean.apply(
+            lambda text: get_bert_embedding(text=text, cls=False, layer_dict=layer_dict)
+        )
+        return df
+
+    if "my_bert_cls_embedding" not in df.columns:
+        df["my_bert_cls_embedding"] = None
+
+    if "my_bert_mean_embedding" not in df.columns:
+        df["my_bert_mean_embedding"] = None
+
+    counter = 0
+
+    df["my_bert_cls_embedding"] = df["my_bert_cls_embedding"].astype(object)
+    df["my_bert_mean_embedding"] = df["my_bert_mean_embedding"].astype(object)
+
+    for i in tqdm(range(len(df))):
+
+        if df.at[i, "my_bert_cls_embedding"] is not None:
+            df.at[i, "my_bert_cls_embedding"] = get_bert_embedding(
+                text=df.at[i, "text_clean"], cls=False, layer_dict=layer_dict
+            )
+            counter = counter + 1
+            sleep(0.5)
+
+        if df.at[i, "my_bert_mean_embedding"] is not None:
+            df.at[i, "my_bert_mean_embedding"] = get_bert_embedding(
+                text=df.at[i, "text_clean"], cls=True, layer_dict=layer_dict
+            )
+            counter = counter + 1
+            sleep(0.5)
+
+        if counter % 50 in [0, 1]:
+            df.to_csv(save_path, index=False)
+
+    df.to_csv(save_path, index=False)
+
+    return df
+
+
+def add_concat_embeddings(series: pd.DataFrame) -> pd.Series:
+    other_features = {"aov": series["aov"]} | series["text_clean_sentiment"]
+
+    for type in ["cls", "mean"]:
+        bert_embedding = series[f"my_bert_{type}_embedding"]
+        series[f"my_full_{type}_embedding"] = get_concat_embedding(
+            bert_embedding=bert_embedding, other_features=other_features
+        )
+
+    return series

src/regression/datasets/__init__.py

@@ -0,0 +1,3 @@
+from .FullModelDatasetTorch import FullModelDatasetTorch
+from .DecoderDatasetTorch import DecoderDatasetTorch
+from .RegressionDataset import RegressionDataset, regression_dataset, regression_dataset_s3

src/regression/training_scripts/__init__.py

@@ -0,0 +1,2 @@
+from .train_full_model_PL import train_full_model_PL
+from .train_decoder_PL import train_decoder_PL