Update app.py

app.py

@@ -7,9 +7,6 @@ from copy import deepcopy
 import requests
 import os.path
 from tqdm import tqdm
-import json
-from dataclasses import dataclass
-from typing import Optional, List
 
 # Set environment variables
 os.environ['RWKV_JIT_ON'] = '1'
@@ -22,28 +19,9 @@ MODELS = {
     "0.4B (Larger)": "RWKV-x070-World-0.4B-v2.9-20250107-ctx4096.pth"
 }
 
-# Model configurations
-MODEL_CONFIGS = {
-    "RWKV-x070-World-0.1B-v2.8-20241210-ctx4096.pth": {
-        "n_layer": 12,
-        "n_embd": 768,
-        "ctx_len": 4096
-    },
-    "RWKV-x070-World-0.4B-v2.9-20250107-ctx4096.pth": {
-        "n_layer": 24,
-        "n_embd": 1024,
-        "ctx_len": 4096
-    }
-}
-
-@dataclass
-class ModelArgs:
-    n_layer: int
-    n_embd: int
-    ctx_len: int
-    vocab_size: int = 65536
-    n_head: int = 16  # Number of attention heads
-    n_att: int = 1024  # Attention dimension
+# Download vocab file if not present
+VOCAB_FILE = "rwkv_vocab_v20230424.txt"
+VOCAB_URL = "https://raw.githubusercontent.com/BlinkDL/ChatRWKV/main/v2/rwkv_vocab_v20230424.txt"
 
 def download_file(url, filename):
     """Generic file downloader with progress bar"""
@@ -69,44 +47,28 @@ def download_model(model_name):
         url = f"https://huggingface.co/BlinkDL/rwkv-7-world/resolve/main/{model_name}"
         download_file(url, model_name)
 
-class CustomPipeline(PIPELINE):
-    def __init__(self, model, vocab_file):
-        super().__init__(model, vocab_file)
-        self.model_args = None
-
-    def set_model_args(self, args: ModelArgs):
-        self.model_args = args
+def ensure_vocab():
+    """Ensure vocab file is present"""
+    if not os.path.exists(VOCAB_FILE):
+        download_file(VOCAB_URL, VOCAB_FILE)
 
 class ModelManager:
     def __init__(self):
         self.current_model = None
         self.current_model_name = None
         self.pipeline = None
+        ensure_vocab()
     
     def load_model(self, model_choice):
         model_file = MODELS[model_choice]
         if model_file != self.current_model_name:
             download_model(model_file)
-            
-            # Get model configuration
-            config = MODEL_CONFIGS[model_file]
-            model_args = ModelArgs(
-                n_layer=config['n_layer'],
-                n_embd=config['n_embd'],
-                ctx_len=config['ctx_len']
-            )
-            
-            # Initialize model with args
             self.current_model = RWKV(
                 model=model_file,
                 strategy='cpu fp32'
             )
-            
-            # Initialize custom pipeline
-            self.pipeline = CustomPipeline(self.current_model, "20B_tokenizer.json")
-            self.pipeline.set_model_args(model_args)
+            self.pipeline = PIPELINE(self.current_model, VOCAB_FILE)
             self.current_model_name = model_file
-            
         return self.pipeline
 
 model_manager = ModelManager()
@@ -143,8 +105,7 @@ def generate_response(
             alpha_decay=alpha_decay,
             token_ban=[],
             token_stop=[],
-            chunk_len=256,
-            model_args=pipeline.model_args  # Pass model args to pipeline
+            chunk_len=256
         )
         
         # Generate response