Update README.md

README.md

@@ -1,19 +1,99 @@
 ---
 language:
-- en
+  - en
 tags:
-- pytorch
-- causal-lm
-- pythia
+  - pytorch
+  - causal-lm
+  - pythia
 license: apache-2.0
 datasets:
-- Dahoas/synthetic-instruct-gptj-pairwise
+  - Dahoas/synthetic-instruct-gptj-pairwise
 ---
 
-This model is created by finetuning `EleutherAI/pythia-2.8b-deduped` on the `Dahoas/synthetic-instruct-gptj-pairwise` for 4 epochs.
+This model is created by finetuning [`EleutherAI/pythia-2.8b-deduped`](https://huggingface.co/EleutherAI/pythia-2.8b-deduped) on the [`Dahoas/synthetic-instruct-gptj-pairwise`](https://huggingface.co/datasets/Dahoas/synthetic-instruct-gptj-pairwise).
 
 You can try a [demo](https://cloud.lambdalabs.com/demos/ml/qa-28b-8000) of the model hosted on [Lambda Cloud](https://lambdalabs.com/service/gpu-cloud).
 
-It took 8xA100 80GB five hours to train the model. We set `batch_size_per_gpu` to `2` (so global batch size is 8), and learning rate to `0.00001` (with linear decay to zero at the last trainig step). 
+### Model Details
 
-The Weights and Biases record of the training can be found [here](https://wandb.ai/chuanli11/ft-synthetic-instruct-gptj-pairwise-pythia2.8b?workspace=user-chuanli11).
+- Finetuned by: [Lambda](https://lambdalabs.com/)
+- Model type: Transformer-based Language Model
+- Language: English
+- Pre-trained model: [EleutherAI/pythia-2.8b-deduped](https://huggingface.co/EleutherAI/pythia-2.8b-deduped)
+- Dataset: [Dahoas/synthetic-instruct-gptj-pairwise](https://huggingface.co/datasets/Dahoas/synthetic-instruct-gptj-pairwise)
+- Library: [transformers](https://huggingface.co/docs/transformers/index)
+- License: Apache 2.0
+
+### Prerequisites
+
+Running inference with the model takes ~7GB of GPU memory.
+
+### Quick Start
+
+```
+import torch
+
+from transformers import AutoTokenizer, pipeline, StoppingCriteria, StoppingCriteriaList
+
+device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
+
+model_name = "lambdalabs/pythia-2.8b-deduped-synthetic-instruct"
+max_new_tokens = 2048
+stop_token = "<|stop|>"
+
+
+class KeywordsStoppingCriteria(StoppingCriteria):
+    def __init__(self, keywords_ids: list):
+        self.keywords = keywords_ids
+
+    def __call__(
+        self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs
+    ) -> bool:
+        if input_ids[0][-1] in self.keywords:
+            return True
+        return False
+
+
+tokenizer = AutoTokenizer.from_pretrained(
+    model_name,
+)
+tokenizer.pad_token = tokenizer.eos_token
+tokenizer.add_tokens([stop_token])
+
+stop_ids = [tokenizer.encode(w)[0] for w in [stop_token]]
+stop_criteria = KeywordsStoppingCriteria(stop_ids)
+
+generator = pipeline(
+    "text-generation",
+    model=model_name,
+    device=device,
+    max_new_tokens=max_new_tokens,
+    torch_dtype=torch.float16,
+    stopping_criteria=StoppingCriteriaList([stop_criteria]),
+)
+
+example = "How can I make an omelette."
+text = "Question: {}\nAnswer:".format(example)
+
+result = generator(
+    text,
+    num_return_sequences=1,
+)
+
+output = result[0]["generated_text"]
+
+print(output)
+```
+
+Output:
+
+```
+Question: How can I make an omelette.
+Answer:To make an omelette, start by cracking two eggs into a bowl and whisking them together. Add a splash of milk and a pinch of salt and pepper. Heat a non-stick pan over medium-high heat and add a tablespoon of butter. Once the butter has melted, pour in the egg mixture. As the eggs set, use a spatula to lift the edges and let the uncooked egg run underneath. When the eggs are cooked through and no visible liquid egg remains, top with your desired fillings and fold the omelette in half before sliding it onto a plate.<|stop|>
+```
+
+### Training
+
+The model was trained on the [`Dahoas/synthetic-instruct-gptj-pairwise`](https://huggingface.co/datasets/Dahoas/synthetic-instruct-gptj-pairwise). We split the original dataset into the train (first 32000 examples) and validation (the remaining 1144 examples) subsets.
+
+We finetune the model for 4 epoches. This took 8xA100 80GB 5 hours, where we set `batch_size_per_gpu` to `2` (so global batch size is 16), and learning rate to `0.00001` (with linear decay to zero at the last trainig step). You can find a Weights and Biases record [here](https://wandb.ai/chuanli11/ft-synthetic-instruct-gptj-pairwise-pythia2.8b?workspace=user-chuanli11).