fix ON_COLAB

app.py

@@ -7,18 +7,18 @@ from huggingface_hub import hf_hub_download
 # nvmlInit()
 # gpu_h = nvmlDeviceGetHandleByIndex(0)
 ctx_limit = 4096
-desc = f'''链接：<a href='https://github.com/BlinkDL/ChatRWKV' target="_blank" style="margin:0 0.5em">ChatRWKV</a><a href='https://github.com/BlinkDL/RWKV-LM' target="_blank" style="margin:0 0.5em">RWKV-LM</a><a href="https://pypi.org/project/rwkv/" target="_blank" style="margin:0 0.5em">RWKV pip package</a><a href="https://zhuanlan.zhihu.com/p/618011122" target="_blank" style="margin:0 0.5em">知乎教程</a>
+desc = f'''链接：<a href='https://colab.research.google.com/drive/1J1gLMMMA8GbD9JuQt6OKmwCTl9mWU0bb?usp=sharing'>太慢了？用Colab自己部署吧</a> <br /> <a href='https://github.com/BlinkDL/ChatRWKV' target="_blank" style="margin:0 0.5em">ChatRWKV</a><a href='https://github.com/BlinkDL/RWKV-LM' target="_blank" style="margin:0 0.5em">RWKV-LM</a><a href="https://pypi.org/project/rwkv/" target="_blank" style="margin:0 0.5em">RWKV pip package</a><a href="https://zhuanlan.zhihu.com/p/618011122" target="_blank" style="margin:0 0.5em">知乎教程</a>
 '''
 
 parser = argparse.ArgumentParser(prog = 'ChatGal RWKV')
 parser.add_argument('--share',action='store_true') 
 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")
-if os.environ['ON_COLAB'] == '1':
+if 'ON_COLAB' in os.environ and os.environ['ON_COLAB'] == '1':
+    os.environ["RWKV_JIT_ON"] = '0'
     os.environ["RWKV_CUDA_ON"] = '1' # if '1' then use CUDA kernel for seq mode (much faster)
     model = RWKV(model=model_path, strategy='cuda bf16')
 else:
@@ -88,8 +88,8 @@ examples = [
 
 女招待: 是吗。那真是太好了
 
-我因为撰稿的需要，而造访了这间位于信州山间的温泉宿驿。""", 200, 0.7, 1.0, 0.05, 0.05],
-    ["翡翠: 欢迎回来，志贵少爷。", 200, 0.7, 1.0, 0.05, 0.05],
+我因为撰稿的需要，而造访了这间位于信州山间的温泉宿驿。""", 200, 0.7, 1.0, 0.0, 0.0],
+    ["翡翠: 欢迎回来，志贵少爷。", 200, 0.7, 1.0, 0.0, 0.0],
     ["""莲华: 你的目的，就是这个万华镜吧？
 
 莲华拿出了万华镜。
@@ -105,7 +105,7 @@ examples = [
 
 深见: 请让我好好看看……
 
-我刚想把手伸过去，莲华就一下子把它收了回去。""", 200, 0.7, 1.0, 0.05, 0.05],
+我刚想把手伸过去，莲华就一下子把它收了回去。""", 200, 0.7, 1.0, 0.00, 0.0],
     ["""嘉祥: 偶尔来一次也不错。
 
 我坐到客厅的沙发上，拍了拍自己的大腿。
@@ -122,20 +122,20 @@ examples = [
 
 我摸摸各自占据住我左右两腿的两颗猫头。
 
-嘉祥: 开心归开心，拜托你们俩别一直乱动啊，很危险的。""", 200, 0.7, 1.0, 0.05, 0.05],
+嘉祥: 开心归开心，拜托你们俩别一直乱动啊，很危险的。""", 200, 0.7, 1.0, 0.0, 0.0],
 ]
 
 iface = gr.Interface(
     fn=infer,
-    description=f'''这是纯网文模型，去除了英文和代码能力，但写小白文更强。<b>请点击例子（在页面底部）</b>，可编辑内容。这里只看输入的最后约1200字，请写好，标点规范，无错别字，否则电脑会模仿你的错误。<b>为避免占用资源，每次生成限制长度。可将输出内容复制到输入，然后继续生成</b>。推荐提高temp改善文采，降低topp改善逻辑，提高两个penalty避免重复，具体幅度请自己实验。{desc}''',
+    description=f'''这是GalGame剧本续写模型（实验性质，不保证效果）。<b>请点击例子（在页面底部）</b>，可编辑内容。这里只看输入的最后约1200字，请写好，标点规范，无错别字，否则电脑会模仿你的错误。<b>为避免占用资源，每次生成限制长度。可将输出内容复制到输入，然后继续生成</b>。推荐提高temp改善文采，降低topp改善逻辑，提高两个penalty避免重复，具体幅度请自己实验。{desc}''',
     allow_flagging="never",
     inputs=[
         gr.Textbox(lines=10, label="Prompt 输入的前文", value="通过基因改造，修真"),  # prompt
         gr.Slider(10, 200, step=10, value=200, label="token_count 每次生成的长度"),  # token_count
         gr.Slider(0.2, 2.0, step=0.1, value=0.7, label="temperature 默认0.7，高则变化丰富，低则保守求稳"),  # temperature
         gr.Slider(0.0, 1.0, step=0.05, value=1.0, label="top_p 默认1.0，高则标新立异，低则循规蹈矩"),  # top_p
-        gr.Slider(0.0, 1.0, step=0.1, value=0.05, label="presencePenalty 默认0.05，避免写过的类似字"),  # presencePenalty
-        gr.Slider(0.0, 1.0, step=0.1, value=0.05, label="countPenalty 默认0.05，额外避免写过多次的类似字"),  # countPenalty
+        gr.Slider(0.0, 1.0, step=0.1, value=0.0, label="presencePenalty 默认0.0，避免写过的类似字"),  # presencePenalty
+        gr.Slider(0.0, 1.0, step=0.1, value=0.0, label="countPenalty 默认0.0，额外避免写过多次的类似字"),  # countPenalty
     ],
     outputs=gr.Textbox(label="Output 输出的续写", lines=28),
     examples=examples,