Adds Flan-T5 seq2seq training

elise/src/configs/__init__.py

@@ -0,0 +1 @@
+from .train_t5 import T5TrainingConfig

elise/src/configs/train_t5.py

@@ -0,0 +1,16 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class T5TrainingConfig:
+    """Training configs for T5 finetuing"""
+
+    train_batch_size: int = 32
+    eval_batch_size: int = 32
+    epochs: int = 10
+    max_length: int = 512
+    learning_rate: float = 3e-4
+    num_warmup_steps: int = 200
+    mixed_precision: str = "bf16"
+    gradient_accumulation_steps: int = 4
+    output_dir: str = "FlanT5_MIT_ner"

elise/src/data/__init__.py

@@ -0,0 +1,5 @@
+"""
+Contians datasets and their connectors for model training
+"""
+
+from .mit_seq2seq_dataset import MITRestaurants, get_default_transforms

elise/src/data/mit_seq2seq_dataset.py

@@ -0,0 +1,122 @@
+"""
+seq2seq models datasets
+
+Classes:
+    MITRestaurants: tner/mit_restaurant dataset to seq2seq
+
+Functions: 
+    get_default_transforms: default transforms for mit dataset
+"""
+import datasets
+
+
+class MITRestaurants:
+    """
+    tner/mit_restaurants for seq2seq
+
+    Atrributes
+    ----------
+    ner_tags: ner tags and ids of mit restaurant
+    dataset: hf dataset
+    transforms: transforms to apply
+    """
+
+    ner_tags = {
+        "O": 0,
+        "B-Rating": 1,
+        "I-Rating": 2,
+        "B-Amenity": 3,
+        "I-Amenity": 4,
+        "B-Location": 5,
+        "I-Location": 6,
+        "B-Restaurant_Name": 7,
+        "I-Restaurant_Name": 8,
+        "B-Price": 9,
+        "B-Hours": 10,
+        "I-Hours": 11,
+        "B-Dish": 12,
+        "I-Dish": 13,
+        "B-Cuisine": 14,
+        "I-Price": 15,
+        "I-Cuisine": 16,
+    }
+
+    def __init__(self, dataset: datasets.DatasetDict, transforms=None):
+        """
+        Constructs mit datasets
+
+        Parameters:
+            dataset: huggingface mit dataset
+            transforms: dataset transform functions
+        """
+        self.dataset = dataset
+        self.transforms = transforms
+
+    def hf_training(self):
+        """
+        Returns dataset for huggingface training ecosystem
+        """
+        dataset_ = self.dataset
+        if self.transforms:
+            for transfrom in self.transforms:
+                dataset_ = dataset_.map(transfrom)
+        return dataset_
+
+    def set_transforms(self, transforms):
+        """
+        Set tranfroms fn
+
+        Parameters:
+            transforms: transforms functions
+        """
+        if self.transforms:
+            self.transforms += transforms
+        else:
+            self.transforms = transforms
+        return self
+
+    @classmethod
+    def from_hf(cls, hf_path: str):
+        """
+        Constructs dataset from huggingface
+
+        Parameters:
+            hf_path: path to dataset hf repo
+        """
+        return cls(datasets.load_dataset(hf_path))
+
+
+def get_default_transforms():
+    label_names = {v: k for k, v in MITRestaurants.ner_tags.items()}
+
+    def decode_tags(tags, words):
+        dict_out = {}
+        word_ = ""
+        for tag, word in zip(tags[::-1], words[::-1]):
+            if tag == 0:
+                continue
+            word_ = word + " " + word_
+            if label_names[tag].startswith("B"):
+                tag_name = label_names[tag][2:]
+                word_ = word_.strip()
+                if tag_name not in dict_out:
+                    dict_out[tag_name] = [word_]
+                else:
+                    dict_out[tag_name].append(word_)
+                word_ = ""
+        return dict_out
+
+    def format_to_text(decoded):
+        text = ""
+        for key, value in decoded.items():
+            text += f"{key}: {', '.join(value)}\n"
+        return text
+
+    def generate_seq2seq_data(example):
+        decoded = decode_tags(example["tags"], example["tokens"])
+        return {
+            "tokens": " ".join(example["tokens"]),
+            "labels": format_to_text(decoded),
+        }
+
+    return [generate_seq2seq_data]

elise/src/train_t5_seq2seq.py

@@ -0,0 +1,183 @@
+import torch
+import evaluate
+import datasets
+from torch.utils.data import DataLoader
+from datasets import load_dataset
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from dataclasses import asdict
+
+from transformers import DataCollatorForSeq2Seq
+from accelerate import Accelerator
+from transformers import get_scheduler
+import numpy as np
+import mlflow
+
+from tqdm.auto import tqdm
+
+from data import MITRestaurants, get_default_transforms
+from utils.logger import get_logger
+from configs import T5TrainingConfig
+
+log = get_logger("Flan_T5")
+log.debug("heloooooooooooo?")
+
+# get dataset
+transforms = get_default_transforms()
+dataset = (
+    MITRestaurants.from_hf("tner/mit_restaurant")
+    .set_transforms(transforms)
+    .hf_training()
+)
+dataset["train"] = datasets.concatenate_datasets([dataset["train"], dataset["test"]])
+# log.info(dataset)
+print(dataset)
+tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
+model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
+
+
+def tokenize(example):
+    tokenized = tokenizer(
+        example["tokens"],
+        text_target=example["labels"],
+        max_length=512,
+        truncation=True,
+    )
+
+    return tokenized
+
+
+tokenized_datasets = dataset.map(
+    tokenize,
+    batched=True,
+    remove_columns=dataset["train"].column_names,
+)
+
+# bleu metric
+metric = evaluate.load("sacrebleu")
+
+
+def postprocess(predictions, labels):
+    predictions = predictions.cpu().numpy()
+    labels = labels.cpu().numpy()
+
+    decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)
+
+    # Replace -100 in the labels as we can't decode them.
+    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
+    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
+
+    # Some simple post-processing
+    decoded_preds = [pred.strip() for pred in decoded_preds]
+    decoded_labels = [[label.strip()] for label in decoded_labels]
+    return decoded_preds, decoded_labels
+
+
+config = T5TrainingConfig()
+
+# data collator
+data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
+
+# data loaders
+tokenized_datasets.set_format("torch")
+train_dataloader = DataLoader(
+    tokenized_datasets["train"],
+    shuffle=True,
+    collate_fn=data_collator,
+    batch_size=config.train_batch_size,
+)
+eval_dataloader = DataLoader(
+    tokenized_datasets["validation"],
+    collate_fn=data_collator,
+    batch_size=config.eval_batch_size,
+)
+
+# optimizer
+optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate)
+num_update_steps_per_epoch = len(train_dataloader)
+num_training_steps = config.epochs * num_update_steps_per_epoch
+
+lr_scheduler = get_scheduler(
+    "linear",
+    optimizer=optimizer,
+    num_warmup_steps=config.num_warmup_steps,
+    num_training_steps=num_training_steps,
+)
+
+# accelerator
+accelerator = Accelerator(
+    mixed_precision=config.mixed_precision,
+    gradient_accumulation_steps=config.gradient_accumulation_steps,
+)
+model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
+    model, optimizer, train_dataloader, eval_dataloader
+)
+
+progress_bar = tqdm(range(num_training_steps))
+
+
+def train(model, dataset, metric):
+    # log.info("Starting Training")
+    print("Starting Traning")
+    for epoch in range(config.epochs):
+        # Training
+        model.train()
+        for batch in train_dataloader:
+            with accelerator.accumulate(model):
+                outputs = model(**batch)
+                loss = outputs.loss
+                accelerator.backward(loss)
+
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+                progress_bar.update(1)
+
+        # Evaluation
+        model.eval()
+        for batch in tqdm(eval_dataloader):
+            with torch.no_grad():
+                generated_tokens = accelerator.unwrap_model(model).generate(
+                    batch["input_ids"],
+                    attention_mask=batch["attention_mask"],
+                    max_length=128,
+                )
+            labels = batch["labels"]
+
+            # Necessary to pad predictions and labels for being gathered
+            generated_tokens = accelerator.pad_across_processes(
+                generated_tokens, dim=1, pad_index=tokenizer.pad_token_id
+            )
+            labels = accelerator.pad_across_processes(labels, dim=1, pad_index=-100)
+
+            predictions_gathered = accelerator.gather(generated_tokens)
+            labels_gathered = accelerator.gather(labels)
+
+            decoded_preds, decoded_labels = postprocess(
+                predictions_gathered, labels_gathered
+            )
+            metric.add_batch(predictions=decoded_preds, references=decoded_labels)
+
+        results = metric.compute()
+        mlflow.log_metrics({"epoch": epoch, "BLEU score": results["score"]})
+        print(f"epoch {epoch}, BLEU score: {results['score']:.2f}")
+
+        # Save and upload
+        accelerator.wait_for_everyone()
+        unwrapped_model = accelerator.unwrap_model(model)
+        unwrapped_model.save_pretrained(
+            config.output_dir, save_function=accelerator.save
+        )
+        if accelerator.is_main_process:
+            tokenizer.save_pretrained(config.output_dir)
+    # save model with mlflow
+    mlflow.transformers.log_model(
+        transformers_model={"model": unwrapped_model, "tokenizer": tokenizer},
+        task="text2text-generation",
+        artifact_path="seq2seq_model",
+        registered_model_name="FlanT5_MIT"
+    )
+
+mlflow.set_tracking_uri("http://127.0.0.1:5000")
+with mlflow.start_run() as mlflow_run:
+    mlflow.log_params(asdict(config))
+    train(model, tokenized_datasets, metric)