update to new model

app.py

@@ -12,11 +12,12 @@ desc = f'''链接：<a href='https://colab.research.google.com/drive/1J1gLMMMA8G
 
 parser = argparse.ArgumentParser(prog = 'ChatGal RWKV')
 parser.add_argument('--share',action='store_true') 
+parser.add_argument('--ckpt',type=str,default="rwkv-chatgal-v1-3B-ctx4096-epoch10.pth") 
 args = parser.parse_args()
 os.environ["RWKV_JIT_ON"] = '1'
 
 from rwkv.model import RWKV
-model_path = hf_hub_download(repo_id="Synthia/ChatGalRWKV", filename="rwkv-chatgal-v1-3B-ctx4096-epoch2.pth")
+model_path = hf_hub_download(repo_id="Synthia/ChatGalRWKV", filename=args.ckpt)
 if 'ON_COLAB' in os.environ and os.environ['ON_COLAB'] == '1':
     os.environ["RWKV_JIT_ON"] = '0'
     os.environ["RWKV_CUDA_ON"] = '0' # if '1' then use CUDA kernel for seq mode (much faster)