update

tamilatis

@@ -1 +0,0 @@
-Subproject commit b1022a9187d9d47c18b360fc45b7f55d3b40824f

tamilatis/configs/config.yaml

@@ -0,0 +1,7 @@
+defaults:
+  - model: default
+  - dataset: default
+  - training: default
+  - wandb: default
+  - override hydra/job_logging: colorlog
+  - override hydra/hydra_logging: colorlog

tamilatis/configs/dataset/default.yaml

@@ -0,0 +1,6 @@
+train_path : "/content/train_intent.pkl"
+valid_path : "/content/val_intent.pkl"
+test_path : "/content/test_intent.pkl"
+output_dir: "/content/saved_models"
+num_labels: 78
+num_intents: 23

tamilatis/configs/model/default.yaml

@@ -0,0 +1,5 @@
+tokenizer_name: "xlm-roberta-base"
+model_name: "xlm-roberta-base"
+num_labels: 78
+num_intents: 23
+test_model : 

tamilatis/configs/training/default.yaml

@@ -0,0 +1,11 @@
+batch_size: 32
+weight_decay: 0.01
+lr: 1e-4
+max_epochs: 20
+patience: 5
+scheduler: "cosine"
+warmup_steps: 0
+do_train: True
+do_predict: False
+ner_cls_path: /content/ner_cls_rlw.csv
+intent_cls_path: /content/intent_cls_rlw.csv

tamilatis/configs/wandb/default.yaml

@@ -0,0 +1,3 @@
+project_name: "tamilatis"
+group_name: "hard-parameter-sharing-rlw"
+run_name: 

tamilatis/dataset.py

@@ -0,0 +1,120 @@
+import pickle
+
+import torch
+from torch.utils.data import Dataset
+from tqdm import tqdm
+from transformers import AutoTokenizer
+
+
+class BuildDataset:
+    def __init__(self):
+        pass
+
+    def tokenize(self, text):
+        """Splits the text and get offsets"""
+        text = text.strip()
+        tokens = text.split()
+        offsets = []
+        for token in tokens:
+            start_idx = text.find(token)
+            end_idx = start_idx + len(token)
+            offsets.append([start_idx, end_idx])
+        return tokens, offsets
+        
+    def convert_to_boi(self, text, annotations):
+        """Convert Intent Tags to BOI Tags"""
+        tokens, offsets = self.tokenize(text)
+        boi_tags = ["O"] * len(tokens)
+
+        for name, value, [start_idx, end_idx] in annotations:
+            value = value.strip()
+            try:
+                token_span = len(value.split())
+
+                start_token_idx = [
+                    token_idx
+                    for token_idx, (s, e) in enumerate(offsets)
+                    if s == start_idx
+                ][0]
+                end_token_idx = start_token_idx + token_span
+                annotation = [name] + ["I" + name[1:]] * (token_span - 1)
+                boi_tags[start_token_idx:end_token_idx] = annotation
+            except Exception as error:
+                pass
+
+        return list(zip(tokens, boi_tags))
+
+    def build_dataset(self, path):
+        """Build a TOD dataset"""
+        with open(path, "rb") as f:
+            data = pickle.load(f)
+
+        boi_data = []
+        for text, annotation, intent in tqdm(data):
+            boi_item = self.convert_to_boi(text, annotation)
+            is_valid = any([True for token, tag in boi_item if tag != "O"])
+            wrong_intent = intent[0] == "B" or intent[0] == "I"
+
+            if is_valid and not wrong_intent:
+                boi_data.append((boi_item, intent))
+        return boi_data
+
+
+class ATISDataset(Dataset):
+    def __init__(self, data, tokenizer, label_encoder, intent_encoder):
+        self.data = data
+        self.label_encoder = label_encoder
+        self.intent_encoder = intent_encoder
+        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer)
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        tokens = [token for token, annotation in self.data[idx][0]]
+        tags = [tag for token, tag in self.data[idx][0]]
+
+        intent_name = self.data[idx][1]
+        intent_label = self.intent_encoder.transform([intent_name])
+        text = "#".join(tokens)
+
+        encoding = self.tokenizer(
+            tokens,
+            max_length=60,
+            padding="max_length",
+            truncation=True,
+            is_split_into_words=True,
+            return_tensors="pt",
+        )
+
+        input_ids = encoding.input_ids.squeeze(0)
+        attention_mask = encoding.attention_mask.squeeze(0)
+        word_ids = encoding.word_ids()
+
+        tags = self.label_encoder.transform(tags)
+
+        labels = []
+        label_all_tokens = None
+        previous_word_idx = None
+
+        for word_idx in word_ids:
+            if word_idx is None:
+                labels.append(-100)
+            elif word_idx != previous_word_idx:
+                labels.append(tags[word_idx])
+            else:
+                labels.append(tags[word_idx] if label_all_tokens else -100)
+            previous_word_idx = word_idx
+
+        labels = torch.tensor(labels)
+        tags = tags.tolist()
+        tags.extend([-100] * (50 - len(tags)))
+
+        return {
+            "text": text,
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels": labels,
+            "intent": intent_label.item(),
+            "tags": tags,
+        }

tamilatis/main.py

@@ -0,0 +1,180 @@
+import logging
+import os
+import pickle
+
+import wandb
+import hydra
+import pandas as pd
+import torch.nn as nn
+import torch.optim as optim
+from accelerate import Accelerator
+from omegaconf.omegaconf import OmegaConf
+from sklearn.preprocessing import LabelEncoder
+from torch.utils.data import DataLoader
+from transformers import AutoTokenizer, get_scheduler
+
+from dataset import ATISDataset, BuildDataset
+from model import JointATISModel
+from predict import TamilATISPredictor
+from trainer import ATISTrainer
+
+logger = logging.getLogger(__name__)
+
+
+@hydra.main(config_path="./configs", config_name="config")
+def main(cfg):
+
+    os.environ['WANDB_PROJECT'] = cfg.wandb.project_name
+    os.environ['WANDB_RUN_GROUP'] = cfg.wandb.group_name
+    
+    logger.info(OmegaConf.to_yaml(cfg, resolve=True))
+    accelerator = Accelerator()
+    # Get all tags
+    annotations = set()
+    intents = set()
+    count = 0
+
+    logger.info("Building Dataset")
+    data_utils = BuildDataset()
+    train_data = data_utils.build_dataset(cfg.dataset.train_path)
+    valid_data = data_utils.build_dataset(cfg.dataset.valid_path)
+    test_data = data_utils.build_dataset(cfg.dataset.test_path)
+
+    annotations, intents, count = set(), set(), 0
+    for boi_data, intent in train_data:
+        if intent[0] == "B" or intent[0] == "I":
+            count += 1
+        intents.add(intent)
+        for token, annotation in boi_data:
+            annotations.add(annotation)
+
+    for boi_data, intent in valid_data:
+        if intent[0] == "B" or intent[0] == "I":
+            count += 1
+        intents.add(intent)
+        for token, annotation in boi_data:
+            annotations.add(annotation)
+
+    for boi_data, intent in test_data:
+        if intent[0] == "B" or intent[0] == "I":
+            count += 1
+        intents.add(intent)
+        for token, annotation in boi_data:
+            annotations.add(annotation)
+
+    annotations = list(annotations)
+    intents = list(intents)
+
+    # convert string labels to int
+    label_encoder = LabelEncoder()
+    label_encoder.fit(annotations)
+
+    intent_encoder = LabelEncoder()
+    intent_encoder.fit(intents)
+
+    train_ds = ATISDataset(
+        train_data, cfg.model.tokenizer_name, label_encoder, intent_encoder
+    )
+    val_ds = ATISDataset(
+        valid_data, cfg.model.tokenizer_name, label_encoder, intent_encoder
+    )
+    test_ds = ATISDataset(
+        test_data, cfg.model.tokenizer_name, label_encoder, intent_encoder
+    )
+
+    train_dl = DataLoader(train_ds, batch_size=cfg.training.batch_size, pin_memory=True)
+    val_dl = DataLoader(val_ds, batch_size=cfg.training.batch_size * 2, pin_memory=True)
+    test_dl = DataLoader(
+        test_ds, batch_size=cfg.training.batch_size * 2, pin_memory=True
+    )
+    logging.info("DataLoaders are created!")
+
+    model = JointATISModel(
+        cfg.model.model_name, cfg.model.num_labels, cfg.model.num_intents
+    )
+    criterion = nn.CrossEntropyLoss()
+    # Optimizer
+    # Split weights in two groups, one with weight decay and the other not.
+    no_decay = ["bias", "LayerNorm.weight"]
+    optimizer_grouped_parameters = [
+        {
+            "params": [
+                p
+                for n, p in model.named_parameters()
+                if not any(nd in n for nd in no_decay)
+            ],
+            "weight_decay": cfg.training.weight_decay,
+        },
+        {
+            "params": [
+                p
+                for n, p in model.named_parameters()
+                if any(nd in n for nd in no_decay)
+            ],
+            "weight_decay": 0.0,
+        },
+    ]
+    optimizer = optim.AdamW(optimizer_grouped_parameters, lr=cfg.training.lr)
+    nb_train_steps = int(
+        len(train_dl) / cfg.training.batch_size * cfg.training.max_epochs
+    )
+
+    if cfg.training.scheduler is not None:
+      scheduler = get_scheduler(
+        cfg.training.scheduler,
+        optimizer,
+        num_warmup_steps=cfg.training.warmup_steps,
+        num_training_steps=nb_train_steps)
+      # Register the LR scheduler
+      accelerator.register_for_checkpointing(scheduler)
+
+    scheduler = None
+    model, optimizer, train_dl, val_dl = accelerator.prepare(
+        model, optimizer, train_dl, val_dl
+    )
+    
+    run = wandb.init(cfg.wandb.project_name,cfg.wandb.group_name,cfg.wandb.run_name)
+    if cfg.training.do_train:
+        trainer = ATISTrainer(
+            model,
+            optimizer,
+            scheduler,
+            criterion,
+            accelerator,
+            cfg.dataset.output_dir,
+            cfg.dataset.num_labels,
+            cfg.dataset.num_intents,
+            run
+        )
+        best_model, best_loss = trainer.fit(
+            cfg.training.max_epochs, train_dl, val_dl, cfg.training.patience
+        )
+        model_dir = f"{cfg.dataset.output_dir}/model_{best_loss}"
+        if not os.path.exists(model_dir):
+            os.makedirs(model_dir)
+        best_model.save_pretrained(model_dir, push_to_hub=False)
+        logging.info(
+            f"The Best model with validation loss {best_loss} is saved in {model_dir}"
+        )
+    if cfg.training.do_predict:
+        predictor = TamilATISPredictor(
+            model,
+            cfg.model.test_model,
+            cfg.model.tokenizer_name,
+            label_encoder,
+            intent_encoder,
+            cfg.model.num_labels,
+        )
+        outputs, intents = predictor.predict_test_data(test_data)
+        ner_cls_rep, intent_cls_rep = predictor.evaluate(outputs, intents)
+        ner_cls_df = pd.DataFrame(ner_cls_rep).transpose()
+        intent_cls_df = pd.DataFrame(intent_cls_rep).transpose()
+        ner_cls_df.to_csv(cfg.training.ner_cls_path)
+        intent_cls_df.to_csv(cfg.training.intent_cls_path)
+        logging.info(
+            f"Classification reports of intents and slots are saved in {cfg.training.ner_cls_path} and {cfg.training.intent_cls_path}"
+        )
+
+
+if __name__ == "__main__":
+    main()

tamilatis/model.py

@@ -0,0 +1,25 @@
+import torch.nn as nn
+from huggingface_hub import PyTorchModelHubMixin
+from transformers import AutoConfig, AutoModelForTokenClassification
+
+
+class JointATISModel(nn.Module, PyTorchModelHubMixin):
+    def __init__(self, model_name, num_labels, num_intents):
+        super().__init__()
+        self.model = AutoModelForTokenClassification.from_pretrained(
+            model_name, num_labels=num_labels
+        )
+        self.model_config = AutoConfig.from_pretrained(model_name)
+        self.intent_head = nn.Linear(self.model_config.hidden_size, num_intents)
+
+    def forward(self, input_ids, attention_mask, labels):
+        outputs = self.model(
+            input_ids, attention_mask, labels=labels, output_hidden_states=True
+        )
+        pooled_output = outputs["hidden_states"][-1][:, 0, :]
+        intent_logits = self.intent_head(pooled_output)
+        return {
+            "dst_logits": outputs.logits,
+            "intent_loss": intent_logits,
+            "dst_loss": outputs.loss,
+        }

tamilatis/predict.py

@@ -0,0 +1,117 @@
+import numpy as np
+import torch
+from tqdm import tqdm
+from transformers import AutoTokenizer
+
+
+
+class TamilATISPredictor:
+    def __init__(
+        self,
+        model,
+        checkpoint_path,
+        tokenizer,
+        label_encoder,
+        intent_encoder,
+        num_labels,
+    ):
+        self.model = model
+        self.model.load_state_dict(torch.load(checkpoint_path))
+        self.model.eval()
+        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer)
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.num_labels = num_labels
+        self.label_encoder = label_encoder
+        self.intent_encoder = intent_encoder
+
+    def get_predictions(self, text):
+
+        inputs = self.tokenizer(
+            text.split(),
+            is_split_into_words=True,
+            return_offsets_mapping=True,
+            padding="max_length",
+            truncation=True,
+            max_length=60,
+            return_tensors="pt",
+        )
+        ids = inputs["input_ids"].to(self.device)
+        mask = inputs["attention_mask"].to(self.device)
+
+        # forward pass
+        loss_dict = self.model(input_ids=ids, attention_mask=mask, labels=None)
+        slot_logits, intent_logits, slot_loss = (
+            loss_dict["dst_logits"],
+            loss_dict["intent_loss"],
+            loss_dict["dst_loss"],
+        )
+
+        active_logits = slot_logits.view(
+            -1, self.num_labels
+        )  # shape (batch_size * seq_len, num_labels)
+        flattened_predictions = torch.argmax(
+            active_logits, axis=1
+        )  # shape (batch_size*seq_len,) - predictions at the token level
+        tokens = self.tokenizer.convert_ids_to_tokens(ids.squeeze().tolist())
+        token_predictions = self.label_encoder.inverse_transform(
+            [i for i in flattened_predictions.cpu().numpy()]
+        )
+        wp_preds = list(
+            zip(tokens, token_predictions)
+        )  # list of tuples. Each tuple = (wordpiece, prediction)
+
+        slot_prediction = []
+        for token_pred, mapping in zip(
+            wp_preds, inputs["offset_mapping"].squeeze().tolist()
+        ):
+            # only predictions on first word pieces are important
+            if mapping[0] == 0 and mapping[1] != 0 and token_pred[0] != "▁":
+                slot_prediction.append(token_pred[1])
+            else:
+                continue
+        intent_preds = torch.argmax(intent_logits, axis=1)
+        intent_preds = self.intent_encoder.inverse_transform(intent_preds.cpu().numpy())
+
+        return intent_preds, slot_prediction
+
+    def predict_test_data(self, test_data):
+        outputs = []
+        intents = []
+
+        for item, intent in tqdm(test_data):
+            try:
+                tokens = [token for token, tag in item]
+                tags = [tag for token, tag in item]
+                text = " ".join(tokens)
+                intent_preds, slot_preds = self.get_predictions(text)
+                outputs.append((tags, slot_preds))
+                intents.append((intent, intent_preds.item()))
+            except Exception as error:
+                print(error)
+        return outputs, intents
+
+    def evaluate(self, outputs, intents):
+        for output in tqdm(outputs):
+            assert len(output[0]) == len(output[1])
+        y_true = [output[0] for output in outputs]
+        y_pred = [output[1] for output in outputs]
+        from seqeval.metrics import classification_report
+
+        ner_cls_rep = classification_report(y_true, y_pred, output_dict=True)
+        from sklearn.metrics import classification_report
+
+        # Compute metrics for intent
+        y_true = self.intent_encoder.transform(
+            [output[0] for output in intents]
+        ).tolist()
+        y_pred = self.intent_encoder.transform(
+            [output[1] for output in intents]
+        ).tolist()
+
+        target_names = self.intent_encoder.classes_.tolist()
+        target_names = [target_names[idx] for idx in np.unique(y_true + y_pred)]
+        intent_cls_rep = classification_report(
+            y_true, y_pred, target_names=target_names, output_dict=True
+        )
+
+        return ner_cls_rep, intent_cls_rep

tamilatis/trainer.py

@@ -0,0 +1,284 @@
+import logging
+import os
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import wandb
+from torchmetrics.functional import accuracy, f1_score, precision, recall
+from tqdm import tqdm, trange
+
+logger = logging.getLogger(__name__)
+
+
+class ATISTrainer:
+    """A Trainer class consists of utitlity functions for training the model"""
+    def __init__(
+        self,
+        model,
+        optimizer,
+        scheduler,
+        criterion,
+        accelerate,
+        output_dir,
+        num_labels,
+        num_intents,
+        run
+    ):
+        self.model = model
+        self.criterion = criterion
+        self.optimizer = optimizer
+        self.scheduler = scheduler
+        self.accelerator = accelerate
+        self.output_dir = output_dir
+        self.num_labels = num_labels
+        self.num_intents = num_intents
+
+        if not os.path.exists(self.output_dir):
+            os.makedirs(self.output_dir)
+
+        self.run = run
+        logging.info(f"Strating Training, outputs are saved in {self.output_dir}")
+
+    def train_step(self, iterator):
+        training_progress_bar = tqdm(iterator, desc="training")
+        for batch in training_progress_bar:
+            input_ids, attention_mask, labels, intents = (
+                batch["input_ids"],
+                batch["attention_mask"],
+                batch["labels"],
+                batch["intent"],
+            )
+            self.optimizer.zero_grad()
+            loss_dict = self.model(input_ids, attention_mask, labels)
+            slot_logits, intent_logits, slot_loss = (
+                loss_dict["dst_logits"],
+                loss_dict["intent_loss"],
+                loss_dict["dst_loss"],
+            )
+
+            # compute training accuracy for slots
+            flattened_target_labels = batch["labels"].view(
+                -1
+            )  # [batch_size * seq_len, ]
+            active_logits = slot_logits.view(
+                -1, self.num_labels
+            )  # [batch_size* seq_len, num_labels]
+            flattened_preds = torch.argmax(
+                active_logits, axis=-1
+            )  # [batch_size * seq_len,]
+
+            # compute accuracy at active labels
+            active_accuracy = (
+                batch["labels"].view(-1) != -100
+            )  # [batch_size * seq_len, ]
+
+            slot_labels = torch.masked_select(flattened_target_labels, active_accuracy)
+            slot_preds = torch.masked_select(flattened_preds, active_accuracy)
+
+            # compute loss for intents
+            #use rlw
+            intent_loss = self.criterion(intent_logits, batch["intent"])
+            weight = F.softmax(torch.randn(1), dim=-1) # RLW is only this!
+            intent_loss = torch.sum(intent_loss*weight.cuda())
+            intent_preds = torch.argmax(intent_logits, axis=1)
+            train_loss = slot_loss + intent_loss
+            self.accelerator.backward(train_loss)
+            self.optimizer.step()
+
+            if self.scheduler is not None:
+              if not self.accelerator.optimizer_step_was_skipped:
+                self.scheduler.step()
+
+            if self.scheduler is not None:
+              self.scheduler.step()
+
+            intent_acc = accuracy(
+                intent_preds, intents, num_classes=self.num_intents, average="weighted"
+            )
+            intent_f1 = f1_score(
+                intent_preds, intents, num_classes=self.num_intents, average="weighted"
+            )
+            intent_rec = recall(
+                intent_preds, intents, num_classes=self.num_intents, average="weighted"
+            )
+            intent_prec = precision(
+                intent_preds, intents, num_classes=self.num_intents, average="weighted"
+            )
+
+            slot_acc = accuracy(
+                slot_preds, slot_labels, num_classes=self.num_labels, average="weighted"
+            )
+            slot_f1 = f1_score(
+                slot_preds, slot_labels, num_classes=self.num_labels, average="weighted"
+            )
+            slot_rec = recall(
+                slot_preds, slot_labels, num_classes=self.num_labels, average="weighted"
+            )
+            slot_prec = precision(
+                slot_preds, slot_labels, num_classes=self.num_labels, average="weighted"
+            )
+
+            self.run.log(
+                {
+                    "train_loss_step": train_loss.cpu().detach().numpy(),
+                    "train_intent_acc_step": intent_acc,
+                    "train_intent_f1_step": intent_f1,
+                    "train_slot_acc_step": slot_acc,
+                    "train_slot_f1_step": slot_f1,
+                }
+            )
+        # logging.info({"train_loss_step": train_loss, "train_intent_acc_step": intent_acc, "train_intent_f1_step": intent_f1, "train_slot_acc_step": slot_acc, "train_slot_f1_step": slot_f1 })
+
+        return {
+            "train_loss_epoch": train_loss / len(iterator),
+            "train_intent_f1_epoch": intent_f1 / len(iterator),
+            "train_intent_acc_epoch": intent_acc / len(iterator),
+            "train_slot_f1_epoch": slot_f1 / len(iterator),
+            "train_slot_acc_epoch": slot_acc / len(iterator),
+        }
+
+    @torch.no_grad()
+    def eval_step(self, iterator):
+        eval_progress_bar = tqdm(iterator, desc="Evaluating")
+        for batch in eval_progress_bar:
+            input_ids, attention_mask, labels, intents = (
+                batch["input_ids"],
+                batch["attention_mask"],
+                batch["labels"],
+                batch["intent"],
+            )
+            loss_dict = self.model(input_ids, attention_mask, labels)
+            slot_logits, intent_logits, slot_loss = (
+                loss_dict["dst_logits"],
+                loss_dict["intent_loss"],
+                loss_dict["dst_loss"],
+            )
+            # compute training accuracy for slots
+            flattened_target_labels = batch["labels"].view(
+                -1
+            )  # [batch_size * seq_len, ]
+            active_logits = slot_logits.view(
+                -1, self.num_labels
+            )  # [batch_size* seq_len, num_labels]
+            flattened_preds = torch.argmax(
+                active_logits, axis=-1
+            )  # [batch_size * seq_len,]
+
+            # compute accuracy at active labels
+            active_accuracy = (
+                batch["labels"].view(-1) != -100
+            )  # [batch_size * seq_len, ]
+
+            slot_labels = torch.masked_select(flattened_target_labels, active_accuracy)
+            slot_preds = torch.masked_select(flattened_preds, active_accuracy)
+
+            # compute loss for intents
+            intent_loss = self.criterion(intent_logits, batch["intent"])
+            weight = F.softmax(torch.randn(1), dim=-1) # RLW is only this!
+            intent_loss = torch.sum(intent_loss*weight.cuda())
+            
+            intent_preds = torch.argmax(intent_logits, axis=1)
+            eval_loss = slot_loss + intent_loss
+
+            intent_acc = accuracy(
+                intent_preds, intents, num_classes=self.num_intents, average="weighted"
+            )
+            intent_f1 = f1_score(
+                intent_preds, intents, num_classes=self.num_intents, average="weighted"
+            )
+            intent_rec = recall(
+                intent_preds, intents, num_classes=self.num_intents, average="weighted"
+            )
+            intent_prec = precision(
+                intent_preds, intents, num_classes=self.num_intents, average="weighted"
+            )
+
+            slot_acc = accuracy(
+                slot_preds, slot_labels, num_classes=self.num_labels, average="weighted"
+            )
+            slot_f1 = f1_score(
+                slot_preds, slot_labels, num_classes=self.num_labels, average="weighted"
+            )
+            slot_rec = recall(
+                slot_preds, slot_labels, num_classes=self.num_labels, average="weighted"
+            )
+            slot_prec = precision(
+                slot_preds, slot_labels, num_classes=self.num_labels, average="weighted"
+            )
+
+            self.run.log(
+                {
+                    "eval_loss_step": eval_loss,
+                    "eval_intent_acc_step": intent_acc,
+                    "eval_intent_f1_step": intent_f1,
+                    "eval_slot_acc_step": slot_acc,
+                    "eval_slot_f1_step": slot_f1,
+                }
+            )
+
+        return {
+            "eval_loss_epoch": eval_loss / len(iterator),
+            "eval_intent_f1_epoch": intent_f1 / len(iterator),
+            "eval_intent_acc_epoch": intent_acc / len(iterator),
+            "eval_slot_f1_epoch": slot_f1 / len(iterator),
+            "eval_slot_acc_epoch": slot_acc / len(iterator),
+        }
+
+    def fit(self, n_epochs, train_dataloader, eval_dataloader, patience):
+        best_eval_loss = float("inf")
+        pbar = trange(n_epochs)
+
+        for epoch in pbar:
+            train_metrics_dict = self.train_step(train_dataloader)
+            eval_metrics_dict = self.eval_step(eval_dataloader)
+            # access all the values from the dicts
+            train_loss, eval_loss = (
+                train_metrics_dict["train_loss_epoch"],
+                eval_metrics_dict["eval_loss_epoch"],
+            )
+            train_intent_f1, eval_intent_f1 = (
+                train_metrics_dict["train_intent_f1_epoch"],
+                eval_metrics_dict["eval_intent_f1_epoch"],
+            )
+            train_intent_acc, eval_intent_acc = (
+                train_metrics_dict["train_intent_acc_epoch"],
+                eval_metrics_dict["eval_intent_acc_epoch"],
+            )
+            train_slot_f1, eval_slot_f1 = (
+                train_metrics_dict["train_intent_acc_epoch"],
+                eval_metrics_dict["eval_intent_acc_epoch"],
+            )
+            train_slot_acc, eval_slot_acc = (
+                train_metrics_dict["train_slot_acc_epoch"],
+                eval_metrics_dict["eval_slot_acc_epoch"],
+            )
+
+
+            if eval_loss < best_eval_loss:
+                best_model = self.model
+                best_eval_loss = eval_loss
+
+                train_logs = {
+                    "epoch": epoch,
+                    "train_loss": train_loss,
+                    "eval_loss": eval_loss,
+                    "train_intent_acc": train_intent_acc,
+                    "train_intent_f1": train_intent_f1,
+                    "eval_intent_f1": eval_intent_f1,
+                    "eval_intent_acc": eval_intent_acc,
+                    "train_slot_f1": train_slot_f1,
+                    "train_slot_acc": train_slot_acc,
+                    "lr": {self.optimizer.param_groups[0]["lr"]: 0.2},
+                }
+
+                train_logs["patience"] = patience
+                logging.info(train_logs)
+                logging.info(eval_metrics_dict)
+
+                self.accelerator.wait_for_everyone()
+                model = self.accelerator.unwrap_model(self.model)
+                self.accelerator.save_state(self.output_dir)
+                logging.info(f"Checkpoint is saved in {self.output_dir}")
+
+        return best_model, best_eval_loss