Update app.py

app.py

@@ -18,11 +18,12 @@ nvmlInit()
 gpu_h = nvmlDeviceGetHandleByIndex(0)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-ctx_limit = 3000
+ctx_limit = 2500
+gen_limit = 500
 ########################## text rwkv ################################################################
 from rwkv.utils import PIPELINE, PIPELINE_ARGS
 
-title_v6 = "RWKV-x060-World-1B6-v2.1-20240328-ctx4096"
+title_v6 = "RWKV-x060-World-3B-v2.1-20240417-ctx4096"
 model_path_v6 = hf_hub_download(repo_id="BlinkDL/rwkv-6-world", filename=f"{title_v6}.pth")
 model_v6 = RWKV(model=model_path_v6, strategy='cuda fp16')
 pipeline_v6 = PIPELINE(model_v6, "rwkv_vocab_v20230424")
@@ -109,7 +110,7 @@ def evaluate(
     yield out_str.strip()
 
 examples = [
-    ["Assistant: How can we craft an engaging story featuring vampires on Mars? Let's think step by step and provide an expert response.", 700, 1, 0.3, 0, 1],
+    ["Assistant: How can we craft an engaging story featuring vampires on Mars? Let's think step by step and provide an expert response.", 500, 1, 0.3, 0, 1],
     ["Assistant: How can we persuade Elon Musk to follow you on Twitter? Let's think step by step and provide an expert response.", 500, 1, 0.3, 0, 1],
     [generate_prompt("東京で訪れるべき素晴らしい場所とその紹介をいくつか挙げてください。"), 500, 1, 0.3, 0, 1],
     [generate_prompt("Write a story using the following information.", "A man named Alex chops a tree down."), 500, 1, 0.3, 0, 1],
@@ -275,7 +276,7 @@ with gr.Blocks(title=title) as demo:
         with gr.Row():
             with gr.Column():
                 prompt = gr.Textbox(lines=2, label="Prompt", value="Assistant: How can we craft an engaging story featuring vampires on Mars? Let's think step by step and provide an expert response.")
-                token_count = gr.Slider(10, 700, label="Max Tokens", step=10, value=700)
+                token_count = gr.Slider(10, gen_limit, label="Max Tokens", step=10, value=gen_limit)
                 temperature = gr.Slider(0.2, 2.0, label="Temperature", step=0.1, value=1.0)
                 top_p = gr.Slider(0.0, 1.0, label="Top P", step=0.05, value=0.3)
                 presence_penalty = gr.Slider(0.0, 1.0, label="Presence Penalty", step=0.1, value=0)