Safety pipeline prototype
diff --git a/safety/README.md b/safety/README.md
new file mode 100644
index 0000000000..2b6f23ff0f
--- /dev/null
+++ b/safety/README.md
@@ -0,0 +1,47 @@
+# Train & Evaluate Safety models
+
+This is the Open Assistant Safety Folder and contains the following:
+
+- Model training scripts
+- Model infrence scripts
+- Data processing scripts
+
+## Mission Statment
+
+Our mission at LAION-AI OpenAssistant safety team is to create a safety pipeline
+that is not only compatible with the OpenAssistant model and project but can
+also integrate with other systems outside of it. We are dedicated to making this
+pipeline modifiable and robust to accommodate the diverse preferences of our
+users.
+
+We understand that our users come from different backgrounds and use various
+types of hardware. Therefore, we strive to make our safety pipeline accessible
+and able to run on consumer hardware, so everyone can benefit from its
+protective features.
+
+Through our commitment to innovation and collaboration, we will continue to
+provide safety solutions that ensure the well-being of our users and the wider
+community.
+
+## Why create a safety pipeline?
+
+Open source and extendable safety pipelines unfortunately do not exist on the
+same on the same scale as those in ChatGPT and other commerical systems. To
+further research in implementable, accurate, and extendable safety pipelines,
+Open Assistant Safety Team will continue to push models and code to the public.
+Much research has been done in things like toxicity detection and bias
+mitigation in LLMs, however the implementation of such research in systems that
+use language models as conversational agents in production settings has largely
+gone undocumented. Furthermore, safety systems that interact with diverse
+communities of users must be able accommodate user prefrences. This is paramount
+in introducing LLM based systems all over the world. We hope that our work will
+generate more research in this field, and allow others to create safe LLM based
+systems.
+
+## Training
+
+- Set training configuration using `config.yaml`
+
+```python
+python model_training/t5_trainer.py
+```
diff --git a/safety/model_training/__init__.py b/safety/model_training/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/safety/model_training/config/config.yaml b/safety/model_training/config/config.yaml
new file mode 100644
index 0000000000..35cace8444
--- /dev/null
+++ b/safety/model_training/config/config.yaml
@@ -0,0 +1,18 @@
+defaults:
+  - trainer: default
+padding_side: "right"
+truncation_side: "right"
+model: "t5-base"
+epochs: 1
+batch_size: 8
+save_folder: "safetyfiles"
+max_length: 256
+special_tokens:
+  context_token: "<ctx>"
+  sep_token: "<sep>"
+  label_token: "<cls>"
+  rot_token: "<rot>"
+dataset:
+  name: "allenai/prosocial-dialog"
+  train: ["train", "validation"]
+  test: "test"
diff --git a/safety/model_training/config/trainer/default.yaml b/safety/model_training/config/trainer/default.yaml
new file mode 100644
index 0000000000..b13dc0b7e7
--- /dev/null
+++ b/safety/model_training/config/trainer/default.yaml
@@ -0,0 +1,4 @@
+_target_: transformers.TrainingArguments
+output_dir: "."
+per_device_train_batch_size: 5
+fp16: False
diff --git a/safety/model_training/custom_datasets/__init__.py b/safety/model_training/custom_datasets/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/safety/model_training/custom_datasets/rot_dataset.py b/safety/model_training/custom_datasets/rot_dataset.py
new file mode 100644
index 0000000000..171451a595
--- /dev/null
+++ b/safety/model_training/custom_datasets/rot_dataset.py
@@ -0,0 +1,103 @@
+from dataclasses import dataclass
+from typing import Dict, List
+
+import torch
+from datasets import concatenate_datasets
+from torch.utils.data import Dataset
+
+LABEL2ID = {
+    "__casual__": "__casual__",
+    "__needs_caution__": "__needs_caution__",
+    "__needs_intervention__": "__needs_intervention__",
+    "__probably_needs_caution__": "__probably_needs_caution__",
+    "__possibly_needs_caution__": "__possibly_needs_caution__",
+}
+
+
+class SafetyDataset(Dataset):
+
+    """
+    Dataset to train safety model with context and ROT from prosocial-dialog
+    input format : input<ctx>context</s>
+    output format : <cls>safety_label<rot>ROTs</s>
+
+    """
+
+    def __init__(self, dataset, split, tokenizer, max_len=512):
+        super().__init__()
+
+        if isinstance(split, List):
+            self.split = "-".join(split)
+            self.dataset = concatenate_datasets([dataset[sp] for sp in split])
+        else:
+            self.split = split
+            self.dataset = dataset[split]
+
+        self.max_len = max_len
+        self.tokenizer = tokenizer
+        self.label2id = LABEL2ID
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        idx_start = idx
+        end = self.dataset[max(0, idx_start - 1)]["episode_done"]
+        while (not end) and (idx_start > 0):
+            end = self.dataset[max(0, idx_start - 2)]["episode_done"]
+            idx_start -= 1
+        idx_start = max(0, idx_start)
+        context = [
+            f'\nUser: {self.dataset[i]["context"]}\n bot:{self.dataset[i]["response"]}' for i in range(idx_start, idx)
+        ]
+        context = self.tokenizer.sep_token.join(context)
+        rots = self.dataset[idx]["rots"]
+        label = self.label2id[self.dataset[idx]["safety_label"]]
+        input_tokens = self.tokenizer.encode(self.dataset[idx]["context"], add_special_tokens=False)
+        max_len = self.max_len - (len(input_tokens) + 2)
+        context = self.tokenizer.encode(
+            context,
+            add_special_tokens=False,
+            max_length=max_len,
+        )
+        rots = self.tokenizer.sep_token.join(rots)
+        input_ids = input_tokens + [self.tokenizer.context_token_id] + context + [self.tokenizer.eos_token_id]
+        input_ids = input_ids + [self.tokenizer.pad_token_id] * max(0, (self.max_len - len(input_ids)))
+        mask = [1] * len(input_ids) + [self.tokenizer.pad_token_id] * (self.max_len - len(input_ids))
+        target_text = self.tokenizer.label_token + label + self.tokenizer.context_token + rots
+        decoder_ids = self.tokenizer(
+            target_text,
+            add_special_tokens=True,
+            max_length=self.max_len,
+            padding="max_length",
+        )
+
+        return {
+            "input_ids": torch.LongTensor(input_ids),
+            "attention_mask": torch.LongTensor(mask),
+            "decoder_input_ids": torch.LongTensor(decoder_ids["input_ids"]),
+            "decoder_attention_mask": torch.LongTensor(decoder_ids["attention_mask"]),
+        }
+
+
+@dataclass
+class SafetyDataCollator:
+    def __call__(self, batch: List) -> Dict[str, torch.Tensor]:
+        """
+        Take a list of samples from a Dataset and collate them into a batch.
+        Returns:
+        A dictionary of tensors
+        """
+
+        input_ids = torch.stack([example["input_ids"] for example in batch])
+        lm_labels = torch.stack([example["decoder_input_ids"] for example in batch])
+        lm_labels[lm_labels[:, :] == 0] = -100
+        attention_mask = torch.stack([example["attention_mask"] for example in batch])
+        decoder_attention_mask = torch.stack([example["decoder_attention_mask"] for example in batch])
+
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels": lm_labels,
+            "decoder_attention_mask": decoder_attention_mask,
+        }
diff --git a/safety/model_training/t5_trainer.py b/safety/model_training/t5_trainer.py
new file mode 100644
index 0000000000..2776803b61
--- /dev/null
+++ b/safety/model_training/t5_trainer.py
@@ -0,0 +1,50 @@
+import os
+
+import hydra
+from custom_datasets.rot_dataset import SafetyDataCollator, SafetyDataset
+from datasets import load_dataset
+from hydra.utils import instantiate
+from omegaconf import DictConfig, OmegaConf
+from transformers import T5ForConditionalGeneration, T5Tokenizer, Trainer
+from utils import add_special_tokens
+
+
+@hydra.main(version_base=None, config_path="config", config_name="config")
+def train(cfg: DictConfig) -> None:
+    if not os.path.exists(cfg.save_folder):
+        os.mkdir(cfg.save_folder)
+
+    model = T5ForConditionalGeneration.from_pretrained(cfg.model)
+    tokenizer = T5Tokenizer.from_pretrained(
+        cfg.model,
+        padding_side=cfg.padding_side,
+        truncation_side=cfg.truncation_side,
+        model_max_length=model.config.n_positions,
+    )
+    add_special_tokens(cfg.special_tokens, tokenizer, model)
+    training_args = instantiate(cfg.trainer)
+
+    dataset = load_dataset(cfg.dataset.name)
+    train_dataset = SafetyDataset(
+        dataset, split=OmegaConf.to_object(cfg.dataset.train), tokenizer=tokenizer, max_len=cfg.max_length
+    )
+    valid_dataset = SafetyDataset(dataset, split=cfg.dataset.test, tokenizer=tokenizer, max_len=cfg.max_length)
+
+    # Initialize our Trainer
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=valid_dataset,
+        data_collator=SafetyDataCollator(),
+    )
+
+    # Training
+    trainer.train()
+
+    trainer.save_model(os.path.join(cfg.save_folder, f"{cfg.model_name}-model"))
+    tokenizer.save_vocabulary(os.path.join(cfg.save_folder, f"{cfg.model_name}-tokenizer"))
+
+
+if __name__ == "__main__":
+    train()
diff --git a/safety/model_training/utils.py b/safety/model_training/utils.py
new file mode 100644
index 0000000000..8102ea1056
--- /dev/null
+++ b/safety/model_training/utils.py
@@ -0,0 +1,7 @@
+def add_special_tokens(special_tokens, tokenizer, model):
+    for key, value in special_tokens.items():
+        setattr(tokenizer, key, value)
+        tokenizer.add_tokens([value])
+        setattr(tokenizer, key + "_id", tokenizer.encode(value)[0])
+
+    model.resize_token_embeddings(len(tokenizer))
diff --git a/safety/requirements.txt b/safety/requirements.txt
new file mode 100644
index 0000000000..dffd40c36a
--- /dev/null
+++ b/safety/requirements.txt
@@ -0,0 +1,2 @@
+hydra-core==1.3.2
+omegaconf==2.3.0

The safety team has decided to maintain a `safety` folder. This folder will contain

1. Documentation related to safety models.
2. Reproducible training and evaluation code for safety models.
3. Inference pipelines for safety models

@ontocord 
