From d217931470b2e04a527e6b759516e02d9042f284 Mon Sep 17 00:00:00 2001
From: Michael Louis <michaellouis157@gmail.com>
Date: Wed, 12 Jun 2024 08:43:43 -0400
Subject: [PATCH] Added example and added to mkdocs

---
 examples/cerebrium/cerebrium.toml | 27 ++++++++++++++++
 examples/cerebrium/main.py        | 52 +++++++++++++++++++++++++++++++
 mkdocs.yml                        |  1 +
 3 files changed, 80 insertions(+)
 create mode 100644 examples/cerebrium/cerebrium.toml
 create mode 100644 examples/cerebrium/main.py

diff --git a/examples/cerebrium/cerebrium.toml b/examples/cerebrium/cerebrium.toml
new file mode 100644
index 000000000..56f525d85
--- /dev/null
+++ b/examples/cerebrium/cerebrium.toml
@@ -0,0 +1,27 @@
+[cerebrium.deployment]
+name = "cerebrium"
+python_version = "3.11"
+cuda_version = "12"
+include = "[./*, main.py, cerebrium.toml]"
+exclude = "[.*]"
+shell_commands = []
+
+[cerebrium.hardware]
+cpu = 2
+memory = 14.0
+gpu = "AMPERE A10"
+gpu_count = 1
+provider = "aws"
+region = "us-east-1"
+
+[cerebrium.scaling]
+min_replicas = 0
+max_replicas = 5
+cooldown = 60
+
+[cerebrium.dependencies.pip]
+outline = "==0.0.37"
+transformers = "==4.38.2"
+datasets = "==2.18.0"
+accelerate = "==0.27.2"
+
diff --git a/examples/cerebrium/main.py b/examples/cerebrium/main.py
new file mode 100644
index 000000000..0b9bf6c49
--- /dev/null
+++ b/examples/cerebrium/main.py
@@ -0,0 +1,52 @@
+import outlines
+
+model = outlines.models.transformers("mistralai/Mistral-7B-Instruct-v0.2")
+generator = outlines.generate.json(model, schema)
+
+schema = {
+    "title": "Character",
+    "type": "object",
+    "properties": {
+        "name": {
+            "title": "Name",
+            "maxLength": 10,
+            "type": "string"
+        },
+        "age": {
+            "title": "Age",
+            "type": "integer"
+        },
+        "armor": {"$ref": "#/definitions/Armor"},
+        "weapon": {"$ref": "#/definitions/Weapon"},
+        "strength": {
+            "title": "Strength",
+            "type": "integer"
+        }
+    },
+    "required": ["name", "age", "armor", "weapon", "strength"],
+    "definitions": {
+        "Armor": {
+            "title": "Armor",
+            "description": "An enumeration.",
+            "enum": ["leather", "chainmail", "plate"],
+            "type": "string"
+        },
+        "Weapon": {
+            "title": "Weapon",
+            "description": "An enumeration.",
+            "enum": ["sword", "axe", "mace", "spear", "bow", "crossbow"],
+            "type": "string"
+        }
+    }
+}
+
+def generate(
+    prompt: str = "Amiri, a 53 year old warrior woman with a sword and leather armor.",
+):
+
+    character = generator(
+        f"<s>[INST]Give me a character description. Describe {prompt}.[/INST]"
+    )
+
+    print(character)
+    return character
\ No newline at end of file
diff --git a/mkdocs.yml b/mkdocs.yml
index 01e8506ab..ae72ac496 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -101,6 +101,7 @@ nav:
       - Playing chess: cookbook/models_playing_chess.md
       - Run on the cloud:
           - BentoML: cookbook/deploy-using-bentoml.md
+          - Cerebrium: cookbook/deploy-using-cerebrium.md
           - Modal: cookbook/deploy-using-modal.md
   - Docs:
     - reference/index.md