Update main.py

main.py

@@ -1,7 +1,8 @@
 import math
 import random
 import numpy as np
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from rwkv.model import RWKV
+from rwkv.utils import PIPELINE
 
 # Define the Node class for MCTS
 class Node:
@@ -110,15 +111,13 @@ class GameState:
 
 # Initialize the RWKV model and tokenizer
 model_name = "BlinkDL/rwkv-4-raven"
-tokenizer = AutoModelForCausalLM.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+model = RWKV(model=model_name, strategy="cuda fp16")
+pipeline = PIPELINE(model, "20B_tokenizer")
 
 # Generate Chain-of-Thought
 def generate_cot(state):
     input_text = f"Current state: {state}\nWhat is the best move?"
-    inputs = tokenizer(input_text, return_tensors="pt")
-    outputs = model.generate(**inputs, max_length=100, num_return_sequences=1)
-    cot = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    cot = pipeline.generate(input_text, max_tokens=100)
     return cot
 
 # Use CoT in MCTS
@@ -129,7 +128,4 @@ def mcts_with_cot(initial_state):
     return best_state, cot
 
 # Function to be called by Gradio
-def run_mcts_cot(initial_board):
-    initial_state = GameState(initial_board, 1)
-    best_state, cot = mcts_with_cot(initial_state)
-    return str(best_state), cot
+def run_mcts_c