Upload folder using huggingface_hub

neural_engine/Cargo.toml

@@ -9,3 +9,4 @@ shakmaty = "0.26.0" # Fast chess library
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 rand = "0.8.5"
+rand_distr = "0.4.3"

neural_engine/src/bin/test_zobrist.rs

@@ -0,0 +1,8 @@
+use shakmaty::{Chess, Position};
+use shakmaty::zobrist::{Zobrist64, ZobristHash};
+
+fn main() {
+    let board = Chess::default();
+    let hash = board.zobrist_hash::<Zobrist64>(shakmaty::EnPassantMode::Legal);
+    println!("Hash: {}", hash.0);
+}

neural_engine/src/main.rs

@@ -3,7 +3,9 @@ use std::collections::HashMap;
 use std::fs;
 use std::path::Path;
 use shakmaty::{Chess, Position, CastlingMode};
+use shakmaty::zobrist::{Zobrist64, ZobristHash};
 use ort::session::Session;
+use rand_distr::{Dirichlet, Distribution};
 
 #[derive(Clone)]
 struct MCTSNode {
@@ -24,6 +26,15 @@ impl MCTSNode {
     }
 }
 
+fn add_dirichlet_noise(node: &mut MCTSNode, alpha: f32, epsilon: f32) {
+    if node.children.is_empty() { return; }
+    let dirichlet = Dirichlet::new(&vec![alpha; node.children.len()]).unwrap();
+    let noise = dirichlet.sample(&mut rand::thread_rng());
+    for (i, child) in node.children.values_mut().enumerate() {
+        child.prior = (1.0 - epsilon) * child.prior + epsilon * noise[i];
+    }
+}
+
 fn load_vocab() -> (HashMap<String, i64>, HashMap<i64, String>) {
     let data = fs::read_to_string("vocab.json").expect("Unable to read vocab.json. Please run build_engine.bat first to generate it!");
     let vocab: HashMap<String, i64> = serde_json::from_str(&data).expect("JSON format error");
@@ -96,7 +107,7 @@ fn mcts_simulate(
     session: &mut Session,
     vocab: &HashMap<String, i64>,
     inv_vocab: &HashMap<i64, String>,
-    nn_cache: &mut HashMap<String, (HashMap<String, f32>, f32)>,
+    nn_cache: &mut HashMap<u64, (HashMap<String, f32>, f32)>,
 ) -> f32 {
     if board.is_game_over() {
         if board.is_checkmate() {
@@ -123,7 +134,7 @@ fn mcts_simulate(
         }
         
         // Evaluate with ONNX or Cache
-        let cache_key = shakmaty::fen::Fen::from_position(board.clone(), shakmaty::EnPassantMode::Legal).to_string();
+        let cache_key = board.zobrist_hash::<Zobrist64>(shakmaty::EnPassantMode::Legal).0;
         let (policy, value) = if let Some(cached) = nn_cache.get(&cache_key) {
             cached.clone()
         } else {
@@ -201,7 +212,8 @@ fn main() {
     
     let (vocab, inv_vocab) = load_vocab();
 
-    let mut nn_cache = HashMap::new();
+    let mut nn_cache: HashMap<u64, (HashMap<String, f32>, f32)> = HashMap::new();
+    let mut global_root = MCTSNode::new(1.0);
 
     for line in stdin.lock().lines() {
         let line = line.expect("Failed to read line");
@@ -220,6 +232,7 @@ fn main() {
             "position" => {
                 // position startpos moves e2e4 e7e5
                 board = Chess::default();
+                let old_history = history_moves.clone();
                 history_moves.clear();
                 history_moves.push("<bos>".to_string());
                 
@@ -234,6 +247,33 @@ fn main() {
                         }
                     }
                 }
+                
+                // Persist the tree if the history matches
+                let mut matches = true;
+                if old_history.len() <= history_moves.len() {
+                    for i in 0..old_history.len() {
+                        if old_history[i] != history_moves[i] {
+                            matches = false;
+                            break;
+                        }
+                    }
+                } else {
+                    matches = false;
+                }
+                
+                if matches {
+                    for i in old_history.len()..history_moves.len() {
+                        let m = &history_moves[i];
+                        if let Some(child) = global_root.children.remove(m) {
+                            global_root = child;
+                        } else {
+                            global_root = MCTSNode::new(1.0);
+                            break;
+                        }
+                    }
+                } else {
+                    global_root = MCTSNode::new(1.0);
+                }
             }
             "go" => {
                 let mut movetime_ms = 5000;
@@ -290,17 +330,22 @@ fn main() {
                 }
 
                 if let Some(ref mut sess) = session {
-                    let mut root = MCTSNode::new(1.0);
                     let start_time = std::time::Instant::now();
                     // Leave a small buffer (5% or minimum 50ms) to ensure we reply in time
                     let buffer = (movetime_ms as f64 * 0.05) as u128;
                     let safe_movetime = if movetime_ms > buffer { movetime_ms - buffer } else { movetime_ms };
                     let limit = std::time::Duration::from_millis(safe_movetime as u64);
                     
+                    if global_root.children.is_empty() {
+                        mcts_simulate(&mut global_root, board.clone(), &mut history_moves, sess, &vocab, &inv_vocab, &mut nn_cache);
+                    }
+                    add_dirichlet_noise(&mut global_root, 0.3, 0.25);
+                    
                     let mut iter_count = 0;
+                    
                     // Loop continuously until time expires
                     while start_time.elapsed() < limit {
-                        mcts_simulate(&mut root, board.clone(), &mut history_moves, sess, &vocab, &inv_vocab, &mut nn_cache);
+                        mcts_simulate(&mut global_root, board.clone(), &mut history_moves, sess, &vocab, &inv_vocab, &mut nn_cache);
                         iter_count += 1;
                         
                         // Stream UCI info every 100 iterations so !eval works
@@ -308,7 +353,7 @@ fn main() {
                             let mut best_eval = 0.0;
                             let mut best_visits = -1;
                             let mut best_action = String::new();
-                            for (action, child) in &root.children {
+                            for (action, child) in &global_root.children {
                                 if (child.visit_count as i32) > best_visits {
                                     best_visits = child.visit_count as i32;
                                     best_action = action.clone();
@@ -327,7 +372,7 @@ fn main() {
                     let mut best_action = String::new();
                     let mut best_visits = -1;
                     let mut best_eval = 0.0;
-                    for (action, child) in &root.children {
+                    for (action, child) in &global_root.children {
                         if (child.visit_count as i32) > best_visits {
                             best_visits = child.visit_count as i32;
                             best_action = action.clone();