dorkai commited on
Commit
49fc773
1 Parent(s): fa85650

Create run.py

Browse files
Files changed (1) hide show
  1. run.py +48 -0
run.py ADDED
@@ -0,0 +1,48 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #@title Select your model below, then click the play button to start the UI.
2
+ #@markdown Afterwards, just sit tight and wait - the link to the UI should show up after it's done starting up.
3
+
4
+ Model = "Pygmalion 6B" #@param ["Pygmalion 350M", "Pygmalion 1.3B", "Pygmalion 2.7B", "Pygmalion 6B", "Pygmalion 6B Experimental"] {allow-input: true}
5
+
6
+ pretty_model_name_to_hf_name = {
7
+ "Pygmalion 350M": "PygmalionAI/pygmalion-350m",
8
+ "Pygmalion 1.3B": "PygmalionAI/pygmalion-1.3b",
9
+ "Pygmalion 2.7B": "PygmalionAI/pygmalion-2.7b",
10
+ "Pygmalion 6B": "PygmalionAI/pygmalion-6b",
11
+ "Pygmalion 6B Experimental": "PygmalionAI/pygmalion-6b"
12
+ }
13
+
14
+ model_name = pretty_model_name_to_hf_name[Model]
15
+ branch_name = "main" if Model != "Pygmalion 6B Experimental" else "dev"
16
+
17
+ # Copy-pasted from the Kobold notebook. Seems to be necessary for Henk's script
18
+ # to work properly.
19
+ import os
20
+ if not os.path.exists("/content/drive"):
21
+ os.mkdir("/content/drive")
22
+ if not os.path.exists("/content/drive/MyDrive/"):
23
+ os.mkdir("/content/drive/MyDrive/")
24
+
25
+ # Use Henk's easy install code, but pass --init since we'll manually start the
26
+ # server in the background later.
27
+ !wget https://koboldai.org/ckds -O - | bash /dev/stdin --init only
28
+
29
+ # Clone the UI repo and set it up.
30
+ !git clone --depth=1 \
31
+ "https://github.com/PygmalionAI/gradio-ui.git" \
32
+ && cd gradio-ui && pip3 install -r requirements.txt
33
+
34
+ # Start up Kobold in the background.
35
+ # TODO: Figure out a way to keep logs in the foreground so the user knows what's
36
+ # going on.
37
+ print("\n\n\n")
38
+ print("* The model is about to be downloaded and loaded into the GPU.")
39
+ print("* This takes several minutes, sit tight.")
40
+ print("* A link will show up when this step is completed, keep checking back every couple minutes or so.")
41
+ print("\n\n\n")
42
+ os.system(f"cd /content/KoboldAI-Client && python3 aiserver.py --noaimenu --host --port 9090 --model {model_name} --revision {branch_name} --nobreakmodel --lowmem --quiet &")
43
+
44
+ # And start up the UI. It'll wait for Kobold to finish booting up before
45
+ # printing out its URL.
46
+ !python3 gradio-ui/src/app.py \
47
+ --koboldai-url "http://localhost:9090" \
48
+ --share