rixty/KoboldAI-Client
1
1
Fork 0
mirror of https://github.com/KoboldAI/KoboldAI-Client.git synced 2025-06-05 21:59:24 +02:00
Code Issues Projects Releases Wiki Activity

Model saving for colab mode

Browse Source
This commit is contained in:
henk717
2022-03-13 11:22:44 +01:00
parent ccadeabbde
commit 4892556059
2 changed files with 7 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
aiserver.py
View File

@ -738,7 +738,7 @@ parser.add_argument("--nobreakmodel", action='store_true', help="Disables Breakm
parser.add_argument("--unblock", action='store_true', default=False, help="Unblocks the KoboldAI port to be accessible from other machines without optimizing for remote play (It is recommended to use --host instead)") parser.add_argument("--unblock", action='store_true', default=False, help="Unblocks the KoboldAI port to be accessible from other machines without optimizing for remote play (It is recommended to use --host instead)")
parser.add_argument("--quiet", action='store_true', default=False, help="If present will suppress any story related text from showing on the console") parser.add_argument("--quiet", action='store_true', default=False, help="If present will suppress any story related text from showing on the console")
parser.add_argument("--lowmem", action='store_true', help="Extra Low Memory loading for the GPU, slower but memory does not peak to twice the usage") parser.add_argument("--lowmem", action='store_true', help="Extra Low Memory loading for the GPU, slower but memory does not peak to twice the usage")
parser.add_argument("--savemodel", action='store_true', help="Saves the model to the models folder even if --colab is used (Allows you to save models to Google Drive)")
args: argparse.Namespace = None args: argparse.Namespace = None
if(os.environ.get("KOBOLDAI_ARGS") is not None): if(os.environ.get("KOBOLDAI_ARGS") is not None):
import shlex import shlex
@ -1420,7 +1420,7 @@ if(not vars.use_colab_tpu and vars.model not in ["InferKit", "Colab", "OAI", "Go
except Exception as e: except Exception as e:
model = GPTNeoForCausalLM.from_pretrained(vars.model, cache_dir="cache", **lowmem) model = GPTNeoForCausalLM.from_pretrained(vars.model, cache_dir="cache", **lowmem)
if not args.colab: if not args.colab or args.savemodel:
import shutil import shutil
model = model.half() model = model.half()
model.save_pretrained("models/{}".format(vars.model.replace('/', '_'))) model.save_pretrained("models/{}".format(vars.model.replace('/', '_')))

8
colabkobold.sh
View File

@ -2,7 +2,7 @@
# KoboldAI Easy Colab Deployment Script by Henk717 # KoboldAI Easy Colab Deployment Script by Henk717
# read the options # read the options
TEMP=`getopt -o m:i:p:c:d:x:a:l:z:g:t:n:b: --long model:,init:,path:,configname:,download:,aria2:,dloc:xloc:7z:git:tar:ngrok:branch: -- "$@"` TEMP=`getopt -o m:i:p:c:d:x:a:l:z:g:t:n:b:s: --long model:,init:,path:,configname:,download:,aria2:,dloc:xloc:7z:git:tar:ngrok:branch:savemodel: -- "$@"`
eval set -- "$TEMP" eval set -- "$TEMP"
# extract options and their arguments into variables. # extract options and their arguments into variables.
@ -34,6 +34,8 @@ while true ; do
git="$2" ; shift 2 ;; git="$2" ; shift 2 ;;
-b|--branch) -b|--branch)
branch="$2" ; shift 2 ;; branch="$2" ; shift 2 ;;
-s|--savemodel)
savemodel=" --savemodel" ; shift 2 ;;
--) shift ; break ;; --) shift ; break ;;
*) echo "Internal error!" ; exit 1 ;; *) echo "Internal error!" ; exit 1 ;;
esac esac
@ -48,8 +50,8 @@ function launch
exit 0 exit 0
else else
cd /content/KoboldAI-Client cd /content/KoboldAI-Client
echo "Launching KoboldAI with the following options : python3 aiserver.py$model$kmpath$configname$ngrok --colab" echo "Launching KoboldAI with the following options : python3 aiserver.py$model$kmpath$configname$ngrok$savemodel --colab"
python3 aiserver.py$model$kmpath$configname$ngrok --colab python3 aiserver.py$model$kmpath$configname$ngrok$savemodel --colab
exit exit
fi fi
} }