Update app.py

app.py

@@ -7,7 +7,7 @@ from copy import deepcopy
 import requests
 import os.path
 from tqdm import tqdm
-
+import json
 
 # Set environment variables
 os.environ['RWKV_JIT_ON'] = '1'
@@ -20,27 +20,19 @@ MODELS = {
     "0.4B (Larger)": "RWKV-x070-World-0.4B-v2.9-20250107-ctx4096.pth"
 }
 
-# Download tokenizer if not present
-TOKENIZER_FILE = "rwkv_vocab_v20230424.txt"
-TOKENIZER_URL = "https://raw.githubusercontent.com/BlinkDL/ChatRWKV/main/v2/rwkv_vocab_v20230424.txt"
-
-def download_tokenizer():
-    if not os.path.exists(TOKENIZER_FILE):
-        print("Downloading tokenizer...")
-        response = requests.get(TOKENIZER_URL)
-        with open(TOKENIZER_FILE, 'wb') as f:
-            f.write(response.content)
+# Tokenizer settings
+TOKENIZER_FILE = "20B_tokenizer.json"
+TOKENIZER_URL = "https://raw.githubusercontent.com/BlinkDL/ChatRWKV/main/20B_tokenizer.json"
 
-def download_model(model_name):
-    """Download model if not present"""
-    if not os.path.exists(model_name):
-        print(f"Downloading {model_name}...")
-        url = f"https://huggingface.co/BlinkDL/rwkv-7-world/resolve/main/{model_name}"
+def download_file(url, filename):
+    """Generic file downloader with progress bar"""
+    if not os.path.exists(filename):
+        print(f"Downloading {filename}...")
         response = requests.get(url, stream=True)
         total_size = int(response.headers.get('content-length', 0))
         
-        with open(model_name, 'wb') as file, tqdm(
-            desc=model_name,
+        with open(filename, 'wb') as file, tqdm(
+            desc=filename,
             total=total_size,
             unit='iB',
             unit_scale=True,
@@ -50,11 +42,22 @@ def download_model(model_name):
                 size = file.write(data)
                 pbar.update(size)
 
+def download_model(model_name):
+    """Download model if not present"""
+    if not os.path.exists(model_name):
+        url = f"https://huggingface.co/BlinkDL/rwkv-7-world/resolve/main/{model_name}"
+        download_file(url, model_name)
+
+def ensure_tokenizer():
+    """Ensure tokenizer is present"""
+    download_file(TOKENIZER_URL, TOKENIZER_FILE)
+
 class ModelManager:
     def __init__(self):
         self.current_model = None
         self.current_model_name = None
         self.pipeline = None
+        ensure_tokenizer()
     
     def load_model(self, model_name):
         if model_name != self.current_model_name:
@@ -67,7 +70,6 @@ class ModelManager:
             self.current_model_name = model_name
         return self.pipeline
 
-
 model_manager = ModelManager()
 
 def generate_response(
@@ -115,7 +117,8 @@ def generate_response(
         pipeline.generate(ctx, token_count=max_tokens, args=args, callback=callback)
         return response
     except Exception as e:
-        return f"Error: {str(e)}"
+        import traceback
+        return f"Error: {str(e)}\nStack trace: {traceback.format_exc()}"
 
 # Create the Gradio interface
 with gr.Blocks() as demo:
@@ -228,4 +231,4 @@ with gr.Blocks() as demo:
 
 # Launch the demo
 if __name__ == "__main__":
-    demo.launch(ssr_mode=False)
+    demo.launch()