Update main.py

main.py

@@ -1,15 +1,10 @@
-import sys
 import os
 import math
 import random
 import numpy as np
 import gradio as gr
-from rwkv.model import RWKV
 from transformers import AutoTokenizer
-
-# Add RWKV directory to Python path
-rwkv_dir = os.path.join(os.getcwd(), 'RWKV-LM')
-sys.path.append(rwkv_dir)
+from rwkv import RWKV
 
 # Define the Node class for MCTS
 class Node:
@@ -121,7 +116,7 @@ model_name = "BlinkDL/rwkv-4-raven"
 tokenizer = AutoTokenizer.from_pretrained("gpt2")  # Use a tokenizer from a supported model
 
 # Load the RWKV model
-model = RWKV(model_name=model_name, strategy="cuda fp16")
+model = RWKV(model=model_name, strategy="cuda fp16")
 
 # Generate Chain-of-Thought
 def generate_cot(state):