import math
import json
import os
import random
import re
from pageindex_.utils import ChatGPT_API

class HybridMCTSNode:
    def __init__(self, data, parent=None, prior=1.0):
        self.data = data
        self.parent = parent
        self.children = []
        self.visits = 0
        self.value_sum = 0.0
        self.prior = prior  # P(s, a): Probability from LLM Tree Search (Policy)
        
    def is_leaf(self):
        return not self.children

    def expand(self, child_priors):
        """
        Expand the node by creating children with associated priors from the Policy.
        child_priors: List of (child_data, probability)
        """
        for child_data, prob in child_priors:
            self.children.append(HybridMCTSNode(child_data, parent=self, prior=prob))

    def puct_score(self, c_param=1.414):
        """
        Calculate PUCT (Predictor + UCT) score.
        Combines Value (Q) and Policy (P).
        
        Score = Q(s,a) + U(s,a)
        Q(s,a) = Average Value (Exploitation)
        U(s,a) = c_param * P(s,a) * sqrt(Parent_Visits) / (1 + Visits)
        """
        # Q-value (Value Function component)
        if self.visits == 0:
            q_value = 0.0 # Could be initialized to parent's value or 0
        else:
            q_value = self.value_sum / self.visits
            
        # U-value (Policy/LLM Tree Search component + Exploration)
        # P(s,a) is self.prior
        parent_visits = self.parent.visits if self.parent else 1
        u_value = c_param * self.prior * math.sqrt(parent_visits) / (1 + self.visits)
        
        return q_value + u_value

class PageIndexMCTS:
    def __init__(self, name, root_data, query, model=None):
        """
        Initialize PageIndexMCTS.
        root_data: Can be a dictionary (root node) or a list (children of root).
        """
        if isinstance(root_data, list):
             self.root_data = {
                "title": name,  # filename
                "node_id": "root",
                "summary": "Document Root",
                "nodes": root_data
            }
        else:
            self.root_data = root_data

        self.root = HybridMCTSNode(self.root_data, prior=1.0)
        self.query = query
        self.model = model or os.getenv("MODEL_NAME") or "gpt-3.5-turbo"
        
        # Initialize root children using Policy
        self._expand_node(self.root)

    def _get_policy_priors(self, node):
        """
        [Component 1: Policy Network / LLM Tree Search]
        INTERACTION POINT 1:
        Ask LLM: "Given the query, which of these children are most relevant?"
        Returns: List of (child_data, probability)
        """
        if 'nodes' not in node.data or not node.data['nodes']:
            return []
            
        children_data = node.data['nodes']
        
        
        options_text = "\n".join([f"ID {i}: {child.get('title', '')} - {child.get('summary', '')}" 
                                          for i, child in enumerate(children_data)])
                
        prompt = f"""
        Query: "{self.query}"
                
        You are a search engine router. You are at a node in the document tree.
        Below are the child sections. Assign a probability (0.0 to 1.0) to each child based on how likely it contains the answer.
        The sum of probabilities must be 1.0.
            
        Child Sections:
        {options_text}
                
        Respond ONLY with a JSON object mapping IDs to probabilities:
        {{
            "0": 0.1,
            "1": 0.9
        }}
        """
                
        # 2. Call LLM
        response_text = ChatGPT_API(model=self.model, prompt=prompt)
                
        # 3. Parse Result
        try:
            start = response_text.find('{')
            end = response_text.rfind('}') + 1
            json_str = response_text[start:end]
            probs_map = json.loads(json_str)
        except:
            probs_map = {}

        # 4. Map back to children
        priors = []
        for i, child in enumerate(children_data):
            p = float(probs_map.get(str(i), 1.0 / len(children_data)))
            priors.append((child, p))
        return priors

    def _simulate_policy_priors(self, children_data):
        """
        Simulates policy priors using keyword matching and random noise.
        """
        scores = []
        for child in children_data:
            text = (child.get('title', '') + " " + child.get('summary', '')).lower()
            query_terms = self.query.lower().split()
            match_count = sum(1 for term in query_terms if term in text)
            # Add some noise to make it interesting
            score = match_count + random.uniform(0.1, 0.5)
            scores.append(score)
            
        total_score = sum(math.exp(s) for s in scores)
        if total_score == 0:
             probs = [1.0 / len(children_data) for _ in children_data]
        else:
             probs = [math.exp(s) / total_score for s in scores]
        
        return list(zip(children_data, probs))



    def _get_node_value(self, node):
        """
        [Component 2: Value Network / Verification]
        INTERACTION POINT 2:
        Ask LLM: "Does this specific section answer the query? Score 0-1."
        """
        prompt = f"""
        Query: "{self.query}"
          
        Section Title: {node.data.get('title', '')}
        Section Summary: {node.data.get('summary', '')}
                
        Evaluate the relevance of this section to the query.
        Return a score between 0.0 (completely irrelevant) and 1.0 (perfect match).
                
        Respond ONLY with the number.
        """
        response_text = ChatGPT_API(model=self.model, prompt=prompt)
                
        match = re.search(r"0\.\d+|1\.0|0|1", response_text)
        if match:
            return float(match.group())
        return 0.5

    def _simulate_value(self, node):
        """
        Simulates value function based on keyword matching.
        """
        text_content = (node.data.get('title', '') + " " + node.data.get('summary', '')).lower()
        query_terms = self.query.lower().split()
        
        match_count = 0
        for term in query_terms:
            if term in text_content:
                match_count += 1
        
        if len(query_terms) > 0:
            return match_count / len(query_terms)
        return 0.0

    def _expand_node(self, node):
        """Use Policy to create children"""
        child_priors = self._get_policy_priors(node)
        node.expand(child_priors)

    def run_search(self, iterations=50):
        print(f"🔎 MCTS Searching for: '{self.query}'")
        
        for i in range(iterations):
            node = self.root
            
            # 1. Selection
            while not node.is_leaf() and node.children:
                node = max(node.children, key=lambda c: c.puct_score())
            
            # 2. Expansion
            if 'nodes' in node.data and node.data['nodes'] and node.visits > 0:
                self._expand_node(node)
                if node.children:
                    node = node.children[0] 
            
            # 3. Evaluation
            value = self._get_node_value(node)
            
            # 4. Backpropagation
            while node:
                node.visits += 1
                node.value_sum += value
                node = node.parent

    def get_top_results(self, n=3):
        all_nodes = []
        def collect(node):
            if node.visits > 0 and node != self.root:
                all_nodes.append(node)
            for child in node.children:
                collect(child)
        collect(self.root)
        
        # Sort by Visit Count (standard MCTS robust selection) or Value
        sorted_nodes = sorted(all_nodes, key=lambda x: x.visits, reverse=True)
        
        results = []
        for node in sorted_nodes[:n]:
            results.append({
                "title": node.data.get('title', 'Unknown'),
                "summary": node.data.get('summary', 'No summary'),
                "start_index": node.data.get('start_index', 'N/A'),
                "end_index": node.data.get('end_index', 'N/A'),
                "visits": node.visits,
                "value": f"{node.value_sum / node.visits:.3f}" if node.visits > 0 else "0.000",
                "prior": f"{node.prior:.3f}"
            })
        return results


if __name__ == "__main__":
    structure = [
    {
      "title": "Preface",
      "start_index": 1,
      "end_index": 1,
      "node_id": "0000",
      "summary": "The partial document introduces the paper \"Attention Is All You Need,\" which proposes a novel neural network architecture called the Transformer. This model relies exclusively on attention mechanisms, eliminating the need for recurrent or convolutional layers that are common in traditional sequence transduction models. The authors highlight that the Transformer achieves superior performance on machine translation tasks\u2014specifically setting new state-of-the-art BLEU scores on WMT 2014 English-to-German (28.4) and English-to-French (41.8) benchmarks\u2014while being more parallelizable and significantly faster to train. The paper also demonstrates the model\u2019s generalizability by successfully applying it to English constituency parsing with both large and limited datasets. Additionally, the document includes author affiliations, contribution notes, and publication details from the NIPS 2017 conference."
    },
    {
      "title": "Introduction",
      "start_index": 2,
      "end_index": 2,
      "node_id": "0001",
      "summary": "The partial document introduces the Transformer, a novel neural network architecture designed for sequence modeling and transduction tasks such as language modeling and machine translation. It highlights the limitations of existing recurrent neural networks (RNNs), including Long Short-Term Memory (LSTM) and Gated Recurrent Units (GRUs), which rely on sequential computation that hinders parallelization and becomes inefficient for long sequences. While attention mechanisms have improved modeling of long-range dependencies, they are typically used alongside recurrent structures. The Transformer addresses these issues by eliminating recurrence entirely and relying solely on self-attention mechanisms to capture global dependencies between input and output sequences. This design enables greater parallelization, leading to faster training and state-of-the-art performance in tasks like translation. The document also contrasts the Transformer with other approaches like ByteNet and ConvS2S, which use convolutional layers and still face challenges in modeling distant dependencies efficiently. Finally, it outlines the standard encoder-decoder framework adopted by the Transformer, where the encoder produces continuous representations of the input sequence and the decoder generates the output sequence autoregressively."
    },
    {
      "title": "Background",
      "start_index": 2,
      "end_index": 2,
      "node_id": "0002",
      "summary": "The partial document introduces the Transformer, a novel neural network architecture designed for sequence modeling and transduction tasks such as language modeling and machine translation. It highlights the limitations of existing recurrent neural networks (RNNs), including Long Short-Term Memory (LSTM) and Gated Recurrent Units (GRUs), which rely on sequential computation that hinders parallelization and becomes inefficient for long sequences. While attention mechanisms have improved modeling of long-range dependencies, they are typically used alongside recurrent structures. The Transformer addresses these issues by eliminating recurrence entirely and relying solely on self-attention mechanisms to capture global dependencies between input and output sequences. This design enables greater parallelization, leading to faster training and state-of-the-art performance in tasks like translation. The document also contrasts the Transformer with other approaches like ByteNet and ConvS2S, which use convolutional layers and still face challenges in modeling distant dependencies efficiently. Finally, it outlines the standard encoder-decoder framework adopted by the Transformer, where the encoder produces continuous representations of the input sequence and the decoder generates the output sequence autoregressively."
    },
    {
      "title": "Model Architecture",
      "start_index": 2,
      "end_index": 3,
      "nodes": [
        {
          "title": "Encoder and Decoder Stacks",
          "start_index": 3,
          "end_index": 3,
          "node_id": "0004",
          "summary": "The partial document describes the architecture of the Transformer model, focusing on its encoder and decoder components. It explains that both the encoder and decoder consist of six identical layers, each containing sub-layers with multi-head self-attention and position-wise fully connected feed-forward networks. The encoder uses residual connections and layer normalization around each sub-layer, with all layers producing 512-dimensional outputs. The decoder includes an additional sub-layer that performs multi-head attention over the encoder output and incorporates masking in its self-attention mechanism to prevent attending to future positions, ensuring autoregressive prediction. The section also begins to define the attention mechanism as a function that maps queries and key-value pairs to a weighted sum output."
        },
        {
          "title": "Attention",
          "start_index": 3,
          "end_index": 3,
          "nodes": [
            {
              "title": "Scaled Dot-Product Attention",
              "start_index": 4,
              "end_index": 4,
              "node_id": "0006",
              "summary": "The partial document explains two key attention mechanisms used in transformer models: Scaled Dot-Product Attention and Multi-Head Attention. It describes Scaled Dot-Product Attention as a method that computes attention weights by taking the dot product of queries and keys, scaling by the square root of the key dimension (\u221adk) to prevent large magnitudes that could lead to small softmax gradients, and applying softmax to weight the values. The document contrasts this with additive attention, noting that scaled dot-product attention is faster and more space-efficient due to optimized matrix operations. It then introduces Multi-Head Attention, which applies multiple parallel attention functions using linearly projected versions of queries, keys, and values, allowing the model to jointly attend to information from different representation subspaces."
            },
            {
              "title": "Multi-Head Attention",
              "start_index": 4,
              "end_index": 5,
              "node_id": "0007",
              "summary": "The partial document describes key components of the Transformer architecture, focusing on attention mechanisms and related layers. It begins by introducing Scaled Dot-Product Attention, explaining how queries, keys, and values are used to compute attention weights via scaled dot products and softmax, with scaling by \u221adk to mitigate gradient issues for large dk. It contrasts this with additive attention, noting the computational efficiency of dot-product attention. Next, it details Multi-Head Attention, which applies multiple attention heads in parallel using learned projections, enabling the model to capture information from different representation subspaces; specific dimensions and parameters (e.g., h=8, dk=dv=64) are provided. The document then outlines three applications of attention in the Transformer: encoder-decoder attention, encoder self-attention, and masked decoder self-attention (to preserve autoregressive properties). Finally, it briefly covers Position-wise Feed-Forward Networks (two linear layers with ReLU), and Embeddings and Softmax, including weight sharing between embedding and pre-softmax layers and scaling by \u221admodel."
            },
            {
              "title": "Applications of Attention in our Model",
              "start_index": 5,
              "end_index": 5,
              "node_id": "0008",
              "summary": "The partial document describes key components of the Transformer architecture. It explains multi-head attention, detailing how it enables the model to jointly attend to information from different representation subspaces using multiple parallel attention heads (specifically 8 heads with reduced dimensionality for computational efficiency). It outlines three applications of attention in the model: encoder-decoder attention (where decoder queries attend to encoder outputs), encoder self-attention (where each position attends to all positions in the previous encoder layer), and decoder self-attention (with masking to preserve auto-regressive properties by preventing leftward information flow). The document also covers the position-wise feed-forward networks\u2014applied identically at each position, consisting of two linear transformations with a ReLU activation\u2014and notes their parameterization and dimensionalities (dmodel = 512, dff = 2048). Finally, it discusses the use of learned embeddings for input and output tokens, weight sharing between embedding layers and the pre-softmax linear transformation, and scaling of embedding weights by \u221admodel."
            }
          ],
          "node_id": "0005",
          "summary": "The partial document describes the architecture of the Transformer model, focusing on its encoder and decoder components. It explains that both the encoder and decoder consist of six identical layers, each containing sub-layers with multi-head self-attention and position-wise fully connected feed-forward networks. The encoder uses residual connections and layer normalization around each sub-layer, with all layers producing 512-dimensional outputs. The decoder includes an additional sub-layer that performs multi-head attention over the encoder output and incorporates masking in its self-attention mechanism to prevent attending to future positions, ensuring autoregressive prediction. The section also introduces the general concept of attention as a function mapping queries and key-value pairs to a weighted sum output."
        },
        {
          "title": "Position-wise Feed-Forward Networks",
          "start_index": 5,
          "end_index": 5,
          "node_id": "0009",
          "summary": "The partial document describes key components of the Transformer architecture. It explains multi-head attention, detailing how it enables the model to jointly attend to information from different representation subspaces using multiple parallel attention heads (specifically 8 heads with reduced dimensionality for computational efficiency). It outlines three applications of attention in the model: encoder-decoder attention (where decoder queries attend to encoder outputs), encoder self-attention (where each position attends to all positions in the previous encoder layer), and decoder self-attention (with masking to preserve auto-regressive properties by preventing leftward information flow). The document also covers the position-wise feed-forward networks\u2014applied identically at each position, consisting of two linear transformations with a ReLU activation\u2014and notes their parameterization per layer. Finally, it discusses the use of learned embeddings for input and output tokens, shared weight matrices between embedding layers and the pre-softmax linear transformation, and scaling of embedding weights by \u221admodel."
        },
        {
          "title": "Embeddings and Softmax",
          "start_index": 5,
          "end_index": 6,
          "node_id": "0010",
          "summary": "The partial document describes key components and design choices of the Transformer architecture. It explains multi-head attention, detailing how it enables the model to jointly attend to information from different representation subspaces, with specific parameter dimensions and computational efficiency. It outlines three applications of attention in the model: encoder-decoder attention, encoder self-attention, and masked decoder self-attention to preserve autoregressive properties. The document also covers position-wise feed-forward networks applied identically across positions, shared embedding and softmax weight matrices scaled by \u221ad_model, and sinusoidal positional encodings that allow the model to capture sequence order and potentially generalize to longer sequences. Finally, it justifies the use of self-attention over recurrent and convolutional layers by comparing computational complexity, parallelization potential, and path lengths for capturing long-range dependencies, highlighting self-attention\u2019s constant sequential operation count and shorter dependency paths."
        },
        {
          "title": "Positional Encoding",
          "start_index": 6,
          "end_index": 6,
          "node_id": "0011",
          "summary": "The partial document discusses the advantages of self-attention mechanisms in sequence modeling by comparing them to recurrent and convolutional layers across three key criteria: computational complexity per layer, degree of parallelizability (measured by minimum sequential operations), and maximum path length for capturing long-range dependencies. It presents Table 1, which summarizes these metrics for different layer types\u2014standard self-attention, recurrent, convolutional, and restricted self-attention\u2014highlighting self-attention\u2019s constant path length and high parallelism. The text also explains the need for positional encodings in models lacking recurrence or convolution, describing the use of fixed sinusoidal positional encodings that enable the model to learn relative positions and potentially generalize to longer sequences than seen during training."
        }
      ],
      "node_id": "0003",
      "summary": "The partial document introduces the Transformer, a novel neural network architecture for sequence transduction that eliminates recurrence and relies entirely on attention mechanisms. It begins by highlighting limitations of recurrent neural networks (RNNs), such as their sequential computation that hinders parallelization, despite advances like LSTM and GRU. While attention mechanisms have improved modeling of long-range dependencies, they have typically been used alongside RNNs. The Transformer addresses these issues by using self-attention to capture global dependencies between input and output positions in constant time, enabling greater parallelization and faster training. The document contrasts the Transformer with prior approaches like ByteNet and ConvS2S, which use convolutional layers but still face challenges in modeling distant dependencies efficiently. It then outlines the Transformer\u2019s encoder-decoder structure, composed of stacked layers featuring multi-head self-attention and position-wise feed-forward networks, with residual connections and layer normalization. The decoder includes an additional attention sub-layer over the encoder output and uses masking to preserve autoregressive properties. The section concludes with a general definition of attention as a function mapping queries and key-value pairs to weighted output vectors."
    },
    {
      "title": "Why Self-Attention",
      "start_index": 6,
      "end_index": 7,
      "node_id": "0012",
      "summary": "The partial document covers several key aspects of a neural network architecture, particularly focusing on self-attention mechanisms. It begins with a comparison of different layer types\u2014self-attention, recurrent, convolutional, and restricted self-attention\u2014in terms of computational complexity per layer, minimum number of sequential operations, and maximum path length for capturing long-range dependencies (Table 1). The document then explains the use of sinusoidal positional encodings to inject sequence order information into the model, justifying this choice over learned embeddings due to better extrapolation to longer sequences. Section 4 elaborates on why self-attention is preferred, highlighting its advantages in parallelization, computational efficiency (especially when sequence length is less than representation dimension), and shorter dependency paths compared to recurrent and convolutional layers. It also notes potential interpretability benefits from attention patterns. Finally, Section 5 outlines the training setup, including datasets (WMT 2014 English-German and English-French), batching strategy, hardware (8 NVIDIA P100 GPUs), training duration, optimizer (Adam with custom learning rate scheduling), and regularization techniques."
    },
    {
      "title": "Training",
      "start_index": 7,
      "end_index": 7,
      "nodes": [
        {
          "title": "Training Data and Batching",
          "start_index": 7,
          "end_index": 7,
          "node_id": "0014",
          "summary": "The partial document discusses architectural and training considerations for a neural sequence model, likely a transformer. It compares self-attention with convolutional and recurrent approaches, noting that self-attention can be computationally optimized by limiting its range to local neighborhoods, which affects path length. It highlights the efficiency of separable convolutions but argues that self-attention combined with point-wise feed-forward layers offers comparable complexity with added interpretability, as attention heads capture syntactic and semantic structures. The training section details the use of WMT 2014 datasets for English-German and English-French translation, byte-pair and word-piece tokenization, batching by sequence length, hardware setup (8 NVIDIA P100 GPUs), training duration for base and large models, an Adam optimizer with a custom learning rate schedule involving warmup steps, and mentions the use of three regularization techniques (though specifics are cut off)."
        },
        {
          "title": "Hardware and Schedule",
          "start_index": 7,
          "end_index": 7,
          "node_id": "0015",
          "summary": "The partial document discusses architectural and training considerations for a neural sequence model, likely a transformer. It compares self-attention with convolutional and recurrent approaches, noting that self-attention can be computationally optimized by limiting its range to local neighborhoods, which affects path length. It highlights the efficiency of separable convolutions but argues that self-attention combined with point-wise feed-forward layers offers comparable complexity with added interpretability, as attention heads learn syntactic and semantic patterns. The training section details the use of WMT 2014 datasets for English-German and English-French translation, byte-pair and word-piece tokenization, batching by sequence length, hardware setup (8 NVIDIA P100 GPUs), training duration for base and big models, Adam optimizer settings with a custom learning rate schedule involving warmup steps, and mentions the use of three regularization techniques (though specifics are cut off)."
        },
        {
          "title": "Optimizer",
          "start_index": 7,
          "end_index": 7,
          "node_id": "0016",
          "summary": "The partial document discusses architectural and training considerations for a neural sequence model, likely a transformer. It compares self-attention with convolutional and recurrent approaches, noting that self-attention can be computationally optimized by limiting its range to local neighborhoods, which affects path length. It highlights the efficiency of separable convolutions but argues that self-attention combined with point-wise feed-forward layers offers comparable complexity with added interpretability, as attention heads learn syntactic and semantic patterns. The training section details the use of WMT 2014 datasets for English-German and English-French translation, byte-pair and word-piece tokenization, batch construction by sequence length, hardware setup (8 NVIDIA P100 GPUs), training duration for base and large models, the Adam optimizer with a custom learning rate schedule involving warmup steps, and mentions the use of three regularization techniques (though specifics are cut off)."
        },
        {
          "title": "Regularization",
          "start_index": 7,
          "end_index": 8,
          "node_id": "0017",
          "summary": "The partial document discusses architectural and training aspects of the Transformer model, focusing on efficiency, performance, and empirical results in machine translation. It compares self-attention with convolutional and recurrent approaches, noting that self-attention offers better connectivity and computational efficiency, especially for long sequences, and may enhance model interpretability through attention patterns aligned with syntactic and semantic structures. The training section details the use of WMT 2014 English-German and English-French datasets, byte-pair and word-piece tokenization, batching strategies, hardware setup (8 NVIDIA P100 GPUs), training duration, and hyperparameters. The optimizer is Adam with a custom learning rate schedule involving warmup steps. Regularization techniques include residual dropout and label smoothing. Results show that both base and big Transformer models achieve state-of-the-art BLEU scores on English-to-German and English-to-French tasks, outperforming previous models\u2014including ensembles\u2014while requiring significantly less training cost in FLOPs. Model variations are also mentioned as part of ablation studies to assess component importance."
        }
      ],
      "node_id": "0013",
      "summary": "The partial document discusses architectural and training considerations for a neural sequence model, likely a transformer. It compares self-attention with convolutional and recurrent approaches, noting that self-attention can be computationally optimized by limiting its range to local neighborhoods, which affects path length. It highlights the efficiency of separable convolutions but argues that self-attention combined with point-wise feed-forward layers offers comparable complexity with added interpretability, as attention heads capture syntactic and semantic structures. The training section details the use of WMT 2014 datasets for English-German and English-French translation, byte-pair and word-piece tokenization, batching by sequence length, hardware setup (8 NVIDIA P100 GPUs), training duration for base and large models, an Adam optimizer with a custom learning rate schedule involving warmup steps, and mentions the use of three regularization techniques (though specifics are cut off)."
    },
    {
      "title": "Results",
      "start_index": 8,
      "end_index": 8,
      "nodes": [
        {
          "title": "Machine Translation",
          "start_index": 8,
          "end_index": 8,
          "node_id": "0019",
          "summary": "The partial document highlights the superior performance of the Transformer model in machine translation tasks, specifically on the WMT 2014 English-to-German and English-to-French benchmarks. It presents quantitative results from Table 2, showing that both the base and big Transformer models achieve higher BLEU scores than previous state-of-the-art models\u2014including ensembles\u2014while requiring significantly less training computational cost (measured in FLOPs). The text details training configurations such as dropout rates, label smoothing (\u03f5ls = 0.1), checkpoint averaging strategies, beam search parameters (beam size = 4, length penalty \u03b1 = 0.6), and hardware setup (8 P100 GPUs). It also notes that the big Transformer model sets a new state-of-the-art BLEU score of 28.4 for English-to-German and 41.0 (or 41.8 in Table 2) for English-to-French. Additionally, the document outlines plans to assess the impact of different architectural components through model variations on the English-to-German task."
        },
        {
          "title": "Model Variations",
          "start_index": 8,
          "end_index": 9,
          "node_id": "0020",
          "summary": "The partial document highlights the superior performance of the Transformer model in machine translation tasks, specifically on English-to-German and English-to-French benchmarks (WMT 2014 newstest2014). It reports that both the base and big Transformer models achieve state-of-the-art BLEU scores while requiring significantly less training cost (measured in FLOPs) compared to previous models and ensembles. The text details training strategies such as residual dropout (with Pdrop = 0.1 for base models), label smoothing (\u03f5ls = 0.1), checkpoint averaging, and beam search parameters. It also presents ablation studies (Table 3) examining the impact of architectural variations\u2014including attention head count, key/value dimensions, model size, dropout rates, and positional encoding methods\u2014on translation quality using the English-to-German development set. Finally, it briefly introduces experiments applying the Transformer to English constituency parsing, demonstrating its generalizability to structured prediction tasks with limited data."
        },
        {
          "title": "English Constituency Parsing",
          "start_index": 9,
          "end_index": 10,
          "node_id": "0021",
          "summary": "The partial document discusses variations and evaluations of the Transformer architecture. It presents Table 3, which explores architectural modifications\u2014such as changes to attention heads, key/value dimensions, model size, dropout rates, and positional encoding\u2014and their impact on English-to-German translation performance (measured by perplexity and BLEU score on newstest2013). Key findings include: single-head attention underperforms; too many heads or reduced key size harms quality; larger models improve results; dropout mitigates overfitting; and learned positional embeddings perform similarly to sinusoidal ones. The \"big\" model achieves state-of-the-art BLEU. The document then evaluates the Transformer on English constituency parsing (Table 4), showing strong performance even with minimal task-specific tuning\u2014surpassing prior models in both supervised (WSJ-only) and semi-supervised settings. The conclusion highlights the Transformer as the first fully attention-based sequence transduction model, offering faster training and superior results on WMT translation tasks, with plans to extend it to other modalities and improve efficiency for long sequences."
        }
      ],
      "node_id": "0018",
      "summary": "The partial document highlights the superior performance of the Transformer model in machine translation tasks, specifically on the WMT 2014 English-to-German and English-to-French benchmarks. It presents quantitative results from Table 2, showing that both the base and big Transformer models achieve higher BLEU scores than previous state-of-the-art models\u2014including ensembles\u2014while requiring significantly less training computational cost (measured in FLOPs). The text details training configurations such as dropout rates, label smoothing (\u03f5ls = 0.1), checkpoint averaging strategies, beam search parameters (beam size = 4, length penalty \u03b1 = 0.6), and hardware setup (8 P100 GPUs). It also notes that the big Transformer model sets a new state-of-the-art BLEU score of 28.4 for English-to-German and 41.0 (or 41.8 in Table 2) for English-to-French. Additionally, the document outlines plans to assess the importance of different Transformer components through ablation studies on English-to-German translation."
    },
    {
      "title": "Conclusion",
      "start_index": 10,
      "end_index": 15,
      "node_id": "0022",
      "summary": "The partial document presents experimental results demonstrating the Transformer model's strong performance on English constituency parsing using the WSJ Section 23 benchmark. It compares the Transformer (both 4-layer and semi-supervised variants) against prior models, showing it achieves competitive or state-of-the-art F1 scores\u2014surpassing many earlier approaches and performing nearly as well as the best Recurrent Neural Network Grammar. The text highlights that the Transformer excels even without task-specific tuning and outperforms traditional parsers like the BerkeleyParser when trained only on the standard WSJ dataset. Additionally, the document concludes by summarizing the Transformer as the first sequence transduction model based entirely on attention, noting its faster training speed and superior results on machine translation tasks (WMT 2014 English-to-German and English-to-French). The authors express enthusiasm for extending attention-based models to other modalities and improving efficiency for large-scale inputs. Visualizations of attention heads illustrate the model\u2019s ability to capture long-range dependencies and perform linguistic tasks like anaphora resolution."
    }
  ] # page_index return structure


    query = "Scaled Dot-Product Attention"
    hybrid_search = PageIndexMCTS("1706.03762", structure, query)
    hybrid_search.run_search(iterations=50)
    
    print(f"\n🏆 Top Results (Hybrid MCTS) for '{query}':")
    for res in hybrid_search.get_top_results():
        print(f"- {res['title']}")
        print(f"  Pages: {res['start_index']} - {res['end_index']}")
        print(f"  Visits: {res['visits']} | Value (LLM-V): {res['value']} | Prior (LLM-Policy): {res['prior']}")
        print(f"  Summary: {res['summary']}")