Create main.py

main.py

@@ -0,0 +1,149 @@
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import math
+import random
+import numpy as np
+
+# Define the Node class for MCTS
+class Node:
+    def __init__(self, state, parent=None):
+        self.state = state
+        self.parent = parent
+        self.children = []
+        self.visits = 0
+        self.wins = 0
+
+    def is_fully_expanded(self):
+        return len(self.children) > 0
+
+    def best_child(self, c_param=1.4):
+        choices_weights = [
+            (child.wins / child.visits) + c_param * (2 * math.log(self.visits) / child.visits) ** 0.5 for child in self.children
+        ]
+        return self.children[np.argmax(choices_weights)]
+
+    def expand(self, state):
+        new_node = Node(state, self)
+        self.children.append(new_node)
+        return new_node
+
+# Define the MCTS class
+class MCTS:
+    def __init__(self, simulation_limit=1000):
+        self.root = None
+        self.simulation_limit = simulation_limit
+
+    def search(self, initial_state):
+        self.root = Node(initial_state)
+        for _ in range(self.simulation_limit):
+            node = self.tree_policy(self.root)
+            reward = self.default_policy(node.state)
+            self.backpropagate(node, reward)
+        return self.root.best_child(c_param=0).state
+
+    def tree_policy(self, node):
+        while not node.state.is_terminal():
+            if not node.is_fully_expanded():
+                return self.expand(node)
+            else:
+                node = node.best_child()
+        return node
+
+    def expand(self, node):
+        tried_states = [child.state for child in node.children]
+        new_state = node.state.get_random_child_state()
+        while new_state in tried_states:
+            new_state = node.state.get_random_child_state()
+        return node.expand(new_state)
+
+    def default_policy(self, state):
+        while not state.is_terminal():
+            state = state.get_random_child_state()
+        return state.get_reward()
+
+    def backpropagate(self, node, reward):
+        while node is not None:
+            node.visits += 1
+            node.wins += reward
+            node = node.parent
+
+# Define the Game State and Rules
+class GameState:
+    def __init__(self, board, player):
+        self.board = board
+        self.player = player
+
+    def is_terminal(self):
+        return self.check_win() or self.check_draw()
+
+    def check_win(self):
+        for row in self.board:
+            if row.count(row[0]) == len(row) and row[0] != 0:
+                return True
+        for col in range(len(self.board)):
+            if self.board[0][col] == self.board[1][col] == self.board[2][col] and self.board[0][col] != 0:
+                return True
+        if self.board[0][0] == self.board[1][1] == self.board[2][2] and self.board[0][0] != 0:
+            return True
+        if self.board[0][2] == self.board[1][1] == self.board[2][0] and self.board[0][2] != 0:
+            return True
+        return False
+
+    def check_draw(self):
+        return all(self.board[row][col] != 0 for row in range(len(self.board)) for col in range(len(self.board)))
+
+    def get_random_child_state(self):
+        available_moves = [(row, col) for row in range(len(self.board)) for col in range(len(self.board)) if self.board[row][col] == 0]
+        if not available_moves:
+            return self
+        row, col = random.choice(available_moves)
+        new_board = [row.copy() for row in self.board]
+        new_board[row][col] = self.player
+        return GameState(new_board, 3 - self.player)
+
+    def get_reward(self):
+        if self.check_win():
+            return 1 if self.player == 1 else -1
+        return 0
+
+    def __str__(self):
+        return "\n".join(" ".join(str(cell) for cell in row) for row in self.board)
+
+# Initialize the RWKV model and tokenizer
+model_name = "BlinkDL/rwkv-4-raven"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+
+# 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)
+    return cot
+
+# Use CoT in MCTS
+def mcts_with_cot(initial_state):
+    mcts = MCTS(simulation_limit=1000)
+    best_state = mcts.search(initial_state)
+    cot = generate_cot(best_state)
+    return best_state, cot
+
+# Gradio Interface
+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()