diff --git a/model/supervised_finetuning/README.md b/model/supervised_finetuning/README.md
new file mode 100644
index 00000000..e223e1cd
--- /dev/null
+++ b/model/supervised_finetuning/README.md
@@ -0,0 +1,38 @@
+# Train using supervised examples
+
+Requirements
+
+```
+wandb
+evaluate
+datasets
+transformers
+torch
+```
+
+Start training reward model
+
+```bash
+python trainer.py --configs defaults galactica-125
+```
+
+## Dataset
+
+For now we only support webgpt and summary dataset from OpenAI. Once
+open-asisstant dataset are available it will be added here.
+
+## Model
+
+TBD
+
+## Results
+
+Experimental results in wandb
+[here](https://wandb.ai/sanagnos/supervised-finetuning?workspace=user-sanagnos).
+
+## TODOS
+
+- decide on a model
+- add special token to declare prompt and reply. Do nto freeze the weights for
+  these
+- Merge utils etc with reward model
diff --git a/model/supervised_finetuning/configs/config.yaml b/model/supervised_finetuning/configs/config.yaml
new file mode 100644
index 00000000..f7164002
--- /dev/null
+++ b/model/supervised_finetuning/configs/config.yaml
@@ -0,0 +1,37 @@
+defaults:
+  learning_rate: 1e-5
+  gradient_checkpointing: false
+  gradient_accumulation_steps: 32
+  per_device_train_batch_size: 2
+  per_device_eval_batch_size: 2
+  weight_decay: 0.00
+  warmup_steps: 600
+  eval_steps: 200
+  save_steps: 500
+  max_length: 512
+  num_train_epochs: 3
+  logging_steps: 10
+  max_grad_norm: 2.0
+  save_total_limit: 4
+  eval_accumulation_steps:
+  freeze_layer:
+  datasets:
+    - webgpt
+  cache_dir: ~/.cache
+  loss_fn: CrossEntropyLoss
+  eval_size:
+  log_dir: "base"
+
+galactica-125:
+  learning_rate: 5e-5
+  model_name: facebook/galactica-125m
+  weight_decay: 0.01
+  warmup_steps: 600
+  gradient_checkpointing: false
+  gradient_accumulation_steps: 2
+  per_device_train_batch_size: 4
+  per_device_eval_batch_size: 4
+
+debug:
+  eval_steps: 20
+  eval_size: 100
diff --git a/model/supervised_finetuning/custom_datasets/__init__.py b/model/supervised_finetuning/custom_datasets/__init__.py
new file mode 100644
index 00000000..fcab8a56
--- /dev/null
+++ b/model/supervised_finetuning/custom_datasets/__init__.py
@@ -0,0 +1,67 @@
+from datasets import load_dataset
+from sklearn.model_selection import train_test_split
+from torch.utils.data import Dataset, Subset
+
+
+class SquadV2Dataset(Dataset):
+    def __init__(self, cache_dir, split):
+        self.dataset = load_dataset("squad_v2", cache_dir=cache_dir, split=split)
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        data = self.dataset[idx]
+        # dummy return first answer
+        return "".join([data["title"], ". ", data["context"], " " + data["question"]]), data["answers"]["text"][0]
+
+
+class WebGPT(Dataset):
+    def __init__(self) -> None:
+        super().__init__()
+
+        dataset = load_dataset("openai/webgpt_comparisons")
+        questions = {}
+        # using prompt as our index will allows us
+        # to add additional generated prompt later
+        self.index2question = {}
+        for row in dataset["train"]:
+            question = row["question"]["full_text"]
+            if question not in self.index2question:
+                self.index2question[len(self.index2question)] = question
+
+            # only keep the best answer
+            questions[question] = row["answer_0" if row["score_0"] > row["score_1"] else "answer_1"]
+
+        self.questions = questions
+
+    def __len__(self):
+        return len(self.index2question)
+
+    def __getitem__(self, index):
+        question = self.index2question[index]
+        answer = self.questions[question]
+        return [question, answer]
+
+
+def train_val_dataset(dataset, val_split=0.2):
+    train_idx, val_idx = train_test_split(
+        list(range(len(dataset))), test_size=val_split, random_state=666, shuffle=True
+    )
+    return Subset(dataset, train_idx), Subset(dataset, val_idx)
+
+
+def get_one_dataset(conf, dataset_name):
+    dataset_name = dataset_name.lower()
+
+    if dataset_name == "squadv2":
+        raise ValueError("SquadV2 is not diverse enough for generation .. ")
+        train = SquadV2Dataset(conf.cache_dir, "train")
+        eval = SquadV2Dataset(conf.cache_dir, "validation")
+    elif dataset_name == "webgpt":
+        dataset = WebGPT()
+        train, eval = train_val_dataset(dataset, val_split=0.2)
+    else:
+        raise ValueError(f"Unknown dataset {dataset_name}")
+
+    return train, eval
diff --git a/model/supervised_finetuning/custom_datasets/dialogue_collator.py b/model/supervised_finetuning/custom_datasets/dialogue_collator.py
new file mode 100644
index 00000000..17fe1082
--- /dev/null
+++ b/model/supervised_finetuning/custom_datasets/dialogue_collator.py
@@ -0,0 +1,85 @@
+from dataclasses import dataclass
+from typing import Optional, Union
+
+import numpy as np
+import torch
+from torch.nn import functional as F
+from transformers.tokenization_utils_base import PaddingStrategy, PreTrainedTokenizerBase
+
+
+@dataclass
+class DialogueDataCollator:
+    """
+    Expects a list of texts corresponding to a sequence of [question, answer, question, answer, ...] pairs.
+    """
+
+    tokenizer: PreTrainedTokenizerBase
+    padding: Union[bool, str, PaddingStrategy] = True
+    max_length: Optional[int] = None
+    pad_to_multiple_of: Optional[int] = None
+
+    def __call__(self, features):
+        # TODO add special tokens for question and answer here
+        # additional_special_tokens = ['<question>', '<answer>']
+        prompt_tokens = ["Question: ", "Answer: "]
+
+        flatten_messages = []
+        label_masks = []
+
+        for messages in features:
+            assert len(messages) % 2 == 0, "Number of messages must be even"
+            messages = [
+                (prompt_tokens[0] if i % 2 == 0 else "") + x + ((" " + prompt_tokens[1]) if i % 2 == 0 else "")
+                for i, x in enumerate(messages)
+            ]
+
+            # Add a way for the model to terminate generation, reinitialize prompter
+            messages.append(prompt_tokens[0])
+
+            flatten_messages.append(
+                self.tokenizer(
+                    "".join(messages),
+                    truncation=True,
+                    max_length=self.max_length,
+                    return_offsets_mapping=True,
+                )
+            )
+
+            message_change_indices = np.cumsum([len(x) for x in messages[:-1]])
+            # for each token an integer indicating the index of the message it belongs to. Just to create the label mask.
+            # TEXT:             Question: Hello, how are you? Answer: I am fine. Question: What is your name? Answer: My name is John.
+            # MESSAGE_INDICES:  0         0      0   0   0    0       1 1  1     2         2    2  2    2     2       3  3    3  3
+
+            # If no result in next, we are predicting the last termination token(s)
+            message_indices = list(
+                map(
+                    lambda x: next((i for i, val in enumerate(message_change_indices) if val >= x), -2),
+                    list(map(lambda x: x[1], flatten_messages[-1]["offset_mapping"])),
+                )
+            )
+            label_mask = np.roll(list(map(lambda x: x % 2 == 1, message_indices)), -1, -1)
+            try:
+                label_mask[[i for i in range(len(message_indices)) if message_indices[i] == -2][0] - 1] = True
+            except IndexError:
+                # an aftermath of padding
+                pass
+
+            label_masks.append(label_mask)
+            flatten_messages[-1].pop("offset_mapping")
+
+        batch = self.tokenizer.pad(
+            flatten_messages,
+            padding=self.padding,
+            max_length=self.max_length,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            return_tensors="pt",
+        )
+        dim = batch["input_ids"].shape[-1]
+
+        batch["label_masks"] = torch.stack([F.pad(torch.tensor(x), (0, dim - len(x))) for x in label_masks])
+
+        for k in list(batch.keys()):
+            if k not in ["input_ids", "attention_mask", "label_masks"]:
+                batch.pop(k)
+
+        return batch
diff --git a/model/supervised_finetuning/losses.py b/model/supervised_finetuning/losses.py
new file mode 100644
index 00000000..795396b9
--- /dev/null
+++ b/model/supervised_finetuning/losses.py
@@ -0,0 +1,15 @@
+from torch import nn
+
+
+class CrossEntropyLoss(nn.CrossEntropyLoss):
+    def __init__(self, weight=None, size_average=None, ignore_index=-100, reduce=None, reduction="mean"):
+        super(CrossEntropyLoss, self).__init__(weight, size_average, ignore_index, reduce, reduction)
+
+    def forward(self, input, target, mask=None):
+        if mask is not None:
+            mask = mask.view(-1)
+            input = input.view(-1, input.size(-1))
+            target = target.view(-1)
+            input = input[mask]
+            target = target[mask]
+        return super(CrossEntropyLoss, self).forward(input, target)
diff --git a/model/supervised_finetuning/trainer.py b/model/supervised_finetuning/trainer.py
new file mode 100644
index 00000000..b44890df
--- /dev/null
+++ b/model/supervised_finetuning/trainer.py
@@ -0,0 +1,200 @@
+import argparse
+import os
+from dataclasses import dataclass
+from distutils.util import strtobool
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import torch
+from torch import nn
+from torch.utils.data import Dataset
+from transformers import (
+    DataCollator,
+    EvalPrediction,
+    PreTrainedModel,
+    PreTrainedTokenizerBase,
+    Trainer,
+    TrainerCallback,
+    TrainingArguments,
+    get_cosine_schedule_with_warmup,
+)
+from utils import get_dataset, get_loss, get_model, get_tokenizer, read_yamls
+
+os.environ["WANDB_PROJECT"] = "supervised-finetuning"
+
+
+@dataclass
+class CustomTrainingArguments(TrainingArguments):
+    loss_function: str = "CrossEntropyLoss"
+
+
+def compute_metrics(eval_pred):
+    pred_ids = eval_pred.predictions
+    labels = eval_pred.label_ids
+
+    return {"accuracy": (pred_ids[labels > 0] == labels[labels > 0]).mean()}
+
+
+def preprocess_logits_for_metrics(logits, labels):
+    pred_ids = torch.argmax(logits, dim=-1)
+    return pred_ids
+
+
+class SFTTrainer(Trainer):
+    def __init__(
+        self,
+        model: Union[PreTrainedModel, nn.Module] = None,
+        args: TrainingArguments = None,
+        data_collator: Optional[DataCollator] = None,
+        train_dataset: Optional[Dataset] = None,
+        eval_dataset: Optional[Dataset] = None,
+        tokenizer: Optional[PreTrainedTokenizerBase] = None,
+        model_init: Callable[[], PreTrainedModel] = None,
+        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
+        callbacks: Optional[List[TrainerCallback]] = None,
+        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
+        preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = None,
+    ):
+        super().__init__(
+            model,
+            args,
+            data_collator,
+            train_dataset,
+            eval_dataset,
+            tokenizer,
+            model_init,
+            compute_metrics,
+            callbacks,
+            optimizers,
+            preprocess_logits_for_metrics,
+        )
+        self.loss_fct = get_loss(args.loss_function)
+
+    def fetch_scheduler(self):
+        return get_cosine_schedule_with_warmup(
+            self.optimizer,
+            num_warmup_steps=self.args.warmup_steps,
+            num_training_steps=self.num_train_steps,
+            num_cycles=1,
+            last_epoch=-1,
+        )
+
+    def compute_loss(self, model, inputs, return_outputs=False):
+        labels_mask = inputs.pop("label_masks")
+
+        outputs = model(**inputs)
+
+        loss = self.loss_fct(outputs.get("logits"), torch.roll(inputs["input_ids"], -1, -1), mask=labels_mask)
+
+        return (loss, outputs) if return_outputs else loss
+
+    def _compute_loss(self, model, inputs):
+
+        labels_mask = inputs.pop("label_masks")
+
+        inputs = self._prepare_inputs(inputs)
+
+        outputs = model(**inputs)
+
+        logits = outputs.get("logits")
+
+        targets = torch.roll(inputs["input_ids"], -1, -1)
+        loss = self.loss_fct(outputs.get("logits"), targets, mask=labels_mask)
+
+        return loss, logits, targets, labels_mask
+
+    def prediction_step(
+        self,
+        model: nn.Module,
+        inputs: Dict[str, Union[torch.Tensor, Any]],
+        prediction_loss_only: bool,
+        ignore_keys: Optional[List[str]] = None,
+    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
+
+        with torch.no_grad():
+            loss, logits, labels, labels_mask = self._compute_loss(model, inputs)
+            labels[~labels_mask] = -1
+
+        loss = loss.mean().detach()
+
+        if self.args.prediction_loss_only:
+            return (loss, None, None)
+
+        return (loss, logits, labels)
+
+
+def _strtobool(x):
+    return bool(strtobool(x))
+
+
+def argument_parsing(notebook=False, notebook_args=None):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--configs", nargs="+", required=True)
+
+    if notebook:
+        args, remaining = parser.parse_known_args(notebook_args)
+    else:
+        args, remaining = parser.parse_known_args()
+
+    # Config from YAML
+    conf = {}
+    configs = read_yamls("./configs")
+    for name in args.configs:
+        if "," in name:
+            for n in name.split(","):
+                conf.update(configs[n])
+        else:
+            conf.update(configs[name])
+
+    # Override config from command-line
+    parser = argparse.ArgumentParser()
+    for key, value in conf.items():
+        type_ = type(value) if value is not None else str
+        if type_ == bool:
+            type_ = _strtobool
+        parser.add_argument(f"--{key}", type=type_, default=value)
+
+    return parser.parse_args(remaining)
+
+
+if __name__ == "__main__":
+    training_conf = argument_parsing()
+
+    model = get_model(training_conf)
+    tokenizer = get_tokenizer(training_conf)
+
+    train, evals, collate_fn = get_dataset(training_conf, tokenizer)
+
+    args = CustomTrainingArguments(
+        output_dir=f"{training_conf.model_name}-{training_conf.log_dir}-finetuned",
+        num_train_epochs=training_conf.num_train_epochs,
+        warmup_steps=training_conf.warmup_steps,
+        loss_function=training_conf.loss_fn,
+        learning_rate=float(training_conf.learning_rate),
+        fp16=True,
+        gradient_checkpointing=training_conf.gradient_checkpointing,
+        gradient_accumulation_steps=training_conf.gradient_accumulation_steps,
+        per_device_train_batch_size=training_conf.per_device_train_batch_size,
+        per_device_eval_batch_size=training_conf.per_device_eval_batch_size,
+        weight_decay=training_conf.weight_decay,
+        max_grad_norm=training_conf.max_grad_norm,
+        logging_steps=training_conf.logging_steps,
+        save_total_limit=training_conf.save_total_limit,
+        evaluation_strategy="steps",
+        eval_steps=training_conf.eval_steps,
+        save_steps=training_conf.save_steps,
+        eval_accumulation_steps=training_conf.eval_accumulation_steps,
+        report_to="wandb",
+    )
+
+    assert len(evals) > 0
+    trainer = SFTTrainer(
+        model,
+        args,
+        train_dataset=train,
+        eval_dataset=evals,
+        data_collator=collate_fn,
+        tokenizer=tokenizer,
+        compute_metrics=compute_metrics,
+        preprocess_logits_for_metrics=preprocess_logits_for_metrics,
+    )
+    trainer.train()
diff --git a/model/supervised_finetuning/utils.py b/model/supervised_finetuning/utils.py
new file mode 100644
index 00000000..4a451bed
--- /dev/null
+++ b/model/supervised_finetuning/utils.py
@@ -0,0 +1,111 @@
+from pathlib import Path
+
+import yaml
+from custom_datasets import get_one_dataset
+from custom_datasets.dialogue_collator import DialogueDataCollator
+from losses import CrossEntropyLoss
+from sklearn.model_selection import train_test_split
+from torch.utils.data import ConcatDataset, Subset
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+SUPPORTED_MODELS = ["galactica"]
+
+
+def get_tokenizer(conf):
+    tokenizer = AutoTokenizer.from_pretrained(conf.model_name, cache_dir=conf.cache_dir)
+
+    if "galactica" in conf.model_name:
+        tokenizer.add_special_tokens({"pad_token": "<pad>", "eos_token": "</s>"})
+
+    return tokenizer
+
+
+def get_model(conf):
+    if not any([x in conf.model_name for x in SUPPORTED_MODELS]):
+        raise ValueError(
+            f"Model {conf.model_name} not supported. Supported models: {SUPPORTED_MODELS}. "
+            "To include more make sure the masking is dne correctly... (decoder only supported for now)"
+        )
+
+    model = AutoModelForCausalLM.from_pretrained(conf.model_name, cache_dir=conf.cache_dir)
+
+    if conf.freeze_layer:
+        model = freeze_top_n_layers(model, conf.freeze_layer)
+
+    model_parameters = filter(lambda p: p.requires_grad, model.parameters())
+    params = sum([p.numel() for p in model_parameters])
+    print("Number of trainable parameters: {}M".format(int(params / 1e6)))
+
+    return model
+
+
+def get_dataset(conf, tokenizer):
+    train_datasets, evals = [], {}
+
+    for dataset_name in conf.datasets:
+        train, val = get_one_dataset(conf, dataset_name)
+        train_datasets.append(train)
+        evals[dataset_name] = Subset(val, list(range(min(len(val), conf.eval_size)))) if conf.eval_size else val
+
+    train = ConcatDataset(train_datasets)
+
+    collate_fn = DialogueDataCollator(tokenizer, max_length=conf.max_length)
+
+    return train, evals, collate_fn
+
+
+def get_loss(loss):
+    if loss == "CrossEntropyLoss":
+        return CrossEntropyLoss()
+    else:
+        raise ValueError(f"Loss {loss} not supported")
+
+
+def read_yamls(dir):
+    conf = {}
+    no_conf = True
+
+    for config_file in Path(dir).glob("**/*.yaml"):
+        no_conf = False
+        with config_file.open("r") as f:
+            conf.update(yaml.safe_load(f))
+
+    if no_conf:
+        print(f"WARNING: No yaml files found in {dir}")
+
+    return conf
+
+
+def train_val_dataset(dataset, val_split=0.2):
+    train_idx, val_idx = train_test_split(
+        list(range(len(dataset))), test_size=val_split, random_state=666, shuffle=True
+    )
+    return Subset(dataset, train_idx), Subset(dataset, val_idx)
+
+
+def freeze_top_n_layers(model, target_layers):
+    # its possible we can simply detect which module is a ModuleList
+    # and simply freeze the module without doing string parsing
+    for name, param in model.named_parameters():
+        if "embed" in name:
+            param.requires_grad = False
+        elif ".layer" in name or ".h." in name:
+            tokens = name.split(".")
+            layer_ = None
+            for token in tokens:
+                if token.isdigit():
+                    layer_ = int(token)
+                    break
+
+            if layer_ is not None and layer_ < target_layers:
+                # print('freeze ', layer_, name)
+                param.requires_grad = False
+    return model
+
+
+if __name__ == "__main__":
+    from transformers import AutoModelForSequenceClassification
+
+    model = AutoModelForSequenceClassification.from_pretrained("bigscience/bloomz-560m")
+    freeze_top_n_layers(model, 10)
+    print(model.state_dict().keys())