Update app.py

app.py

@@ -1,32 +1,48 @@
-from sagemaker.huggingface import HuggingFaceModel
-import sagemaker
-
-role = sagemaker.get_execution_role()
-# Hub Model configuration. https://huggingface.co/models
-hub = {
-	'HF_MODEL_ID':'PygmalionAI/pygmalion-6b',
-	'HF_TASK':'conversational'
+#@title Select your model below, then click the play button to start the UI.
+#@markdown Afterwards, just sit tight and wait - the link to the UI should show up after it's done starting up.
+
+Model = "Pygmalion 6B" #@param ["Pygmalion 350M", "Pygmalion 1.3B", "Pygmalion 2.7B", "Pygmalion 6B", "Pygmalion 6B Experimental"] {allow-input: true}
+
+pretty_model_name_to_hf_name = {
+    "Pygmalion 350M": "PygmalionAI/pygmalion-350m",
+    "Pygmalion 1.3B": "PygmalionAI/pygmalion-1.3b",
+    "Pygmalion 2.7B": "PygmalionAI/pygmalion-2.7b",
+    "Pygmalion 6B": "PygmalionAI/pygmalion-6b",
+    "Pygmalion 6B Experimental": "PygmalionAI/pygmalion-6b"
 }
 
-# create Hugging Face Model Class
-huggingface_model = HuggingFaceModel(
-	transformers_version='4.17.0',
-	pytorch_version='1.10.2',
-	py_version='py38',
-	env=hub,
-	role=role, 
-)
-
-# deploy model to SageMaker Inference
-predictor = huggingface_model.deploy(
-	initial_instance_count=1, # number of instances
-	instance_type='ml.m5.xlarge' # ec2 instance type
-)
-
-predictor.predict({
-	'inputs': {
-		"past_user_inputs": ["Which movie is the best ?"],
-		"generated_responses": ["It's Die Hard for sure."],
-		"text": "Can you explain why ?"
-	}
-})
+model_name = pretty_model_name_to_hf_name[Model]
+branch_name = "main" if Model != "Pygmalion 6B Experimental" else "dev"
+
+# Copy-pasted from the Kobold notebook. Seems to be necessary for Henk's script
+# to work properly.
+import os
+if not os.path.exists("/content/drive"):
+  os.mkdir("/content/drive")
+if not os.path.exists("/content/drive/MyDrive/"):
+  os.mkdir("/content/drive/MyDrive/")
+
+# Use Henk's easy install code, but pass --init since we'll manually start the
+# server in the background later.
+!wget https://koboldai.org/ckds -O - | bash /dev/stdin --init only
+
+# Clone the UI repo and set it up.
+!git clone --depth=1 \
+  "https://github.com/PygmalionAI/gradio-ui.git" \
+  && cd gradio-ui && pip3 install -r requirements.txt
+
+# Start up Kobold in the background.
+# TODO: Figure out a way to keep logs in the foreground so the user knows what's
+# going on.
+print("\n\n\n")
+print("* The model is about to be downloaded and loaded into the GPU.")
+print("* This takes several minutes, sit tight.")
+print("* A link will show up when this step is completed, keep checking back every couple minutes or so.")
+print("\n\n\n")
+os.system(f"cd /content/KoboldAI-Client && python3 aiserver.py --noaimenu --host --port 80808 --model {model_name} --revision {branch_name} --nobreakmodel --lowmem --quiet &")
+
+# And start up the UI. It'll wait for Kobold to finish booting up before
+# printing out its URL.
+!python3 gradio-ui/src/app.py \
+  --koboldai-url "http://localhost:80808" \
+  --share