Update app.py

app.py

@@ -8,6 +8,8 @@ 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'
@@ -20,9 +22,28 @@ MODELS = {
     "0.4B (Larger)": "RWKV-x070-World-0.4B-v2.9-20250107-ctx4096.pth"
 }
 
-# Tokenizer settings
-TOKENIZER_FILE = "20B_tokenizer.json"
-TOKENIZER_URL = "https://raw.githubusercontent.com/BlinkDL/ChatRWKV/main/20B_tokenizer.json"
+# 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
 
 def download_file(url, filename):
     """Generic file downloader with progress bar"""
@@ -48,26 +69,44 @@ def download_model(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 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
 
 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:
-            download_model(MODELS[model_name])
+    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=MODELS[model_name],
+                model=model_file,
                 strategy='cpu fp32'
             )
-            self.pipeline = PIPELINE(self.current_model, TOKENIZER_FILE)
-            self.current_model_name = model_name
+            
+            # Initialize custom pipeline
+            self.pipeline = CustomPipeline(self.current_model, "20B_tokenizer.json")
+            self.pipeline.set_model_args(model_args)
+            self.current_model_name = model_file
+            
         return self.pipeline
 
 model_manager = ModelManager()
@@ -104,7 +143,8 @@ def generate_response(
             alpha_decay=alpha_decay,
             token_ban=[],
             token_stop=[],
-            chunk_len=256
+            chunk_len=256,
+            model_args=pipeline.model_args  # Pass model args to pipeline
         )
         
         # Generate response