Update app.py

app.py

@@ -1,31 +1,41 @@
 import os
 import gradio as gr
 from llama_cpp import Llama
-from huggingface_hub import hf_hub_download#, login
+from huggingface_hub import hf_hub_download
 import numpy as np
 
-#login(os.getenv("HF_TOKEN")) my bad now its public
-
 model = Llama(
     model_path=hf_hub_download(
-        repo_id=os.environ.get("REPO_ID", "Lyte/QuadConnect2.5-0.5B-v0.0.3b"),
-        filename=os.environ.get("MODEL_FILE", "quadconnect_q8_0.gguf"),
+        repo_id=os.environ.get("REPO_ID", "Lyte/QuadConnect2.5-0.5B-v0.0.4b"),
+        filename=os.environ.get("MODEL_FILE", "unsloth.Q8_0.gguf"),
     )
 )
 
-SYSTEM_PROMPT = """You are a Connect Four player.  Connect Four is played on a 6x7 grid. Given the current board state, predict the next *column* to play in.  Moves are specified by the column letter (a-g). Respond in the following format:
+SYSTEM_PROMPT = """You are an expert Connect Four player. The game is played on a 6x7 grid where pieces fall to the lowest available position in each column due to gravity.
+
+Board representation:
+- The board is described as a list of occupied positions in the format: <column><row>(<piece>).
+- Columns are labeled a-g (from left to right) and rows are numbered 1-6 (with 1 as the bottom row).
+- For example: 'a1(O), a2(X), b1(O)' indicates that cell a1 has an O, a2 has an X, and b1 has an O.
+- An empty board is simply represented by an empty list.
+- To win, you must connect 4 of your pieces horizontally, vertically, or diagonally.
+
+Respond in the following XML format:
 <reasoning>
-Explain your reasoning, considering the board state and potential future moves.
+Explain your reasoning, including:
+- Identifying winning opportunities for yourself.
+- Blocking your opponent's potential wins.
+- Strategic positioning, such as center control and setting up future moves.
 </reasoning>
 <move>
-The column letter (a-g) of your move.
+Indicate the column letter (a-g) where you want to drop your piece.
 </move>
 """
 
 class ConnectFour:
     def __init__(self):
         self.board = np.zeros((6, 7))
-        self.current_player = 1  # 1 for player, 2 for AI
+        self.current_player = 1  # 1 for player (X), 2 for AI (O)
         self.game_over = False
         
     def make_move(self, col):
@@ -81,14 +91,19 @@ class ConnectFour:
                 if self.board[row][col] != 0:
                     col_letter = chr(ord('a') + col)
                     row_num = str(6 - row)  # Convert to 1-based indexing
-                    #player = "X" if self.board[row][col] == 1 else "O"
-                    moves.append(f"{col_letter}{row_num}")#={player}")
-        return ", ".join(moves)
+                    piece = "X" if self.board[row][col] == 1 else "O"
+                    moves.append(f"{col_letter}{row_num}({piece})")
+        return ", ".join(moves) if moves else ""
 
     def parse_ai_move(self, move_str):
-        # Parse move like 'a1', 'b3', etc.
-        col = ord(move_str[0].lower()) - ord('a')
-        return col
+        # Parse move like 'a', 'b', etc.
+        try:
+            col = ord(move_str.strip().lower()) - ord('a')
+            if 0 <= col <= 6:
+                return col
+            return -1
+        except:
+            return -1
 
 def create_interface():
     game = ConnectFour()
@@ -155,7 +170,7 @@ def create_interface():
 
     with gr.Blocks(css=css) as interface:
         gr.Markdown("# 🎮 Connect Four vs AI")
-        gr.Markdown("### This is just a quick prototype for now, and the current model was trained just for 200 steps to test the concept, the reward functions were flawed, update coming soon!")
+        gr.Markdown("### Play against an AI trained to be an expert Connect Four player!")
         
         with gr.Row():
             with gr.Column(scale=2):
@@ -163,22 +178,22 @@ def create_interface():
                 status = gr.Markdown("Your turn! Click a button to drop your piece!", elem_id="ai-status")
                 
                 # Column buttons
-                with gr.Group(elem_id="column-buttons", elem_classes=["fitter"]):
+                with gr.Group(elem_id="column-buttons"):
                     col_buttons = []
                     for i in range(7):
-                        btn = gr.Button(f"⬇️ {i+1}", scale=1)
+                        btn = gr.Button(f"⬇️ {chr(ord('A') + i)}", scale=1)
                         col_buttons.append(btn)
                 
                 # Game board
                 board_display = gr.HTML(render_board(), elem_id="board-display")
                 reset_btn = gr.Button("🔄 New Game", variant="primary")
             
-        with gr.Column(scale=1):
-            # AI reasoning display
-            gr.Markdown("### 🤖 AI's Thoughts")
-            reasoning_display = gr.HTML(
-                value='<div id="ai-reasoning">Waiting for your move...</div>',
-                elem_id="ai-reasoning-container"
+            with gr.Column(scale=1):
+                # AI reasoning display
+                gr.Markdown("### 🤖 AI's Thoughts")
+                reasoning_display = gr.HTML(
+                    value='<div id="ai-reasoning">Waiting for your move...</div>',
+                    elem_id="ai-reasoning-container"
                 )
 
         def handle_move(col):
@@ -211,9 +226,8 @@ def create_interface():
             # AI move
             game.current_player = 2
             board_state = game.board_to_string()
-
-            prompt = f"Current Board: {board_state}. Make a move."
-            #print(prompt)
+            
+            prompt = f"Current board state (you are O, opponent is X):\n{board_state}\n\nMake your move."
             
             # Get AI response
             response = model.create_chat_completion(
@@ -233,13 +247,16 @@ def create_interface():
                 move_str = ai_response.split("<move>")[1].split("</move>")[0].strip()
                 ai_col = game.parse_ai_move(move_str)
                 
+                if ai_col == -1:
+                    raise ValueError("Invalid move format from AI")
+                
                 # Format reasoning for display
                 reasoning_html = f'''
                 <div id="ai-reasoning">
                     <div class="reasoning-box">
                         <p><strong>🤔 Reasoning:</strong></p>
                         <p>{reasoning}</p>
-                        <p><strong>📍 Move chosen:</strong> {move_str}</p>
+                        <p><strong>📍 Move chosen:</strong> Column {move_str.upper()}</p>
                     </div>
                 </div>
                 '''
@@ -320,7 +337,5 @@ def render_board(board=None):
     html += "</div>"
     return html
 
-# Launch the interface
-if __name__ == "__main__":
-    interface = create_interface()
-    interface.launch()
+interface = create_interface()
+interface.launch()