Update main.py

main.py

@@ -1,8 +1,7 @@
-import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
 import math
 import random
 import numpy as np
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # Define the Node class for MCTS
 class Node:
@@ -111,7 +110,7 @@ class GameState:
 
 # Initialize the RWKV model and tokenizer
 model_name = "BlinkDL/rwkv-4-raven"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+tokenizer = AutoModelForCausalLM.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
 # Generate Chain-of-Thought
@@ -129,21 +128,8 @@ def mcts_with_cot(initial_state):
     cot = generate_cot(best_state)
     return best_state, cot
 
-# Gradio Interface
+# 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
-
-# Create the Gradio interface
-initial_board = [[0, 0, 0], [0, 0, 0], [0, 0, 0]]
-iface = gr.Interface(
-    fn=run_mcts_cot,
-    inputs=gr.inputs.JSON(),
-    outputs=["text", "text"],
-    title="RWKV CoT Demo for MCTS",
-    description="This demo uses RWKV to generate Chain-of-Thought reasoning to guide the MCTS algorithm in a Tic-Tac-Toe game."
-)
-
-# Launch the interface
-iface.launch()
+    return str(best_state), cot