Tree of Thoughts

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from dataclasses import dataclass
from typing import List, Optional, Callable
import re

# Load model
model_name = "Qwen/Qwen2.5-1.5B-Instruct"
print(f"Loading {model_name}...")

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, dtype="auto")

device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
model.eval()

print(f"Loaded on {device}")

@dataclass
class ThoughtNode:
    """
    A node in the Tree of Thoughts.
    
    Each node represents a partial solution state.
    """
    thought: str           # The reasoning step at this node
    state: str             # Full state (all thoughts so far)
    score: float = 0.0     # Evaluation score
    depth: int = 0         # How deep in the tree
    parent: Optional['ThoughtNode'] = None
    children: List['ThoughtNode'] = None
    
    def __post_init__(self):
        if self.children is None:
            self.children = []
    
    def get_path(self) -> List[str]:
        """Get the full reasoning path from root to this node."""
        path = []
        node = self
        while node is not None:
            path.append(node.thought)
            node = node.parent
        return list(reversed(path))


def generate_text(prompt: str, max_new_tokens: int = 50, 
                  temperature: float = 0.7) -> str:
    """Generate text continuation."""
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
        )
    
    full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return full_response[len(prompt):].strip()


print("ThoughtNode defined.")

def generate_thoughts(state: str, n_thoughts: int = 3,
                      problem_type: str = "math") -> List[str]:
    """
    Generate multiple candidate next thoughts.
    
    We ask the model to propose several possible next steps.
    
    Args:
        state: Current reasoning state
        n_thoughts: Number of candidates to generate
        problem_type: Type of problem (affects prompting)
    
    Returns:
        List of candidate next thoughts
    """
    prompt = f"""{state}

What are {n_thoughts} different possible next steps? List them as:
1.
2.
3."""
    
    response = generate_text(prompt, max_new_tokens=100, temperature=0.8)
    
    # Parse numbered list
    thoughts = []
    lines = response.split('\n')
    for line in lines:
        # Match lines starting with a number
        match = re.match(r'^\d+\.\s*(.+)', line.strip())
        if match:
            thoughts.append(match.group(1).strip())
    
    # If parsing failed, just split by newlines
    if len(thoughts) < n_thoughts and lines:
        thoughts = [l.strip() for l in lines if l.strip()][:n_thoughts]
    
    return thoughts[:n_thoughts]


# Test thought generation
test_state = """Problem: Use 2, 3, 4, 6 with +, -, ×, ÷ to make 24.
Current: I have numbers 2, 3, 4, 6."""

print("Generating candidate thoughts...\n")
thoughts = generate_thoughts(test_state, n_thoughts=3)
for i, t in enumerate(thoughts, 1):
    print(f"{i}. {t}")

def evaluate_thought(state: str, thought: str,
                     evaluation_prompt: str = None) -> float:
    """
    Have the model evaluate a thought on a 0-1 scale.
    
    We ask the model to rate the thought as:
    - sure (1.0): This will definitely work
    - likely (0.7): This looks promising
    - possible (0.4): Maybe, not sure
    - unlikely (0.1): Probably won't work
    
    Args:
        state: Current problem state
        thought: The candidate thought to evaluate
        evaluation_prompt: Custom evaluation prompt
    
    Returns:
        Score from 0.0 to 1.0
    """
    if evaluation_prompt is None:
        evaluation_prompt = f"""{state}

Proposed next step: {thought}

Evaluate this step. Is it a good approach to solving the problem?
Rate as: sure / likely / possible / unlikely

Rating:"""
    
    response = generate_text(evaluation_prompt, max_new_tokens=10, temperature=0.3)
    response_lower = response.lower().strip()
    
    # Parse the rating
    if 'sure' in response_lower:
        return 1.0
    elif 'likely' in response_lower:
        return 0.7
    elif 'possible' in response_lower:
        return 0.4
    elif 'unlikely' in response_lower or 'impossible' in response_lower:
        return 0.1
    else:
        # Default to middle score if we can't parse
        return 0.5


# Test evaluation
state = "Problem: Find a word that starts with 'A' and means 'happy'."
thoughts = ["The word 'Awesome' might work.", "The word 'Sad' could be it.", "I'll think about 'Afraid'."]

print("Evaluating thoughts:\n")
for thought in thoughts:
    score = evaluate_thought(state, thought)
    print(f"  {thought}")
    print(f"  → Score: {score:.2f}\n")

class TreeOfThoughtsBFS:
    """
    Tree of Thoughts with Breadth-First Search.
    
    Explores all options at each level, keeps top-k,
    then goes deeper.
    """
    
    def __init__(self, model, tokenizer, device="cuda",
                 n_candidates: int = 3, beam_width: int = 2, max_depth: int = 3):
        self.model = model
        self.tokenizer = tokenizer
        self.device = device
        self.n_candidates = n_candidates  # How many thoughts to generate per node
        self.beam_width = beam_width      # How many nodes to keep at each level
        self.max_depth = max_depth        # How deep to search
    
    def solve(self, problem: str, is_solution: Callable[[str], bool] = None) -> dict:
        """
        Solve a problem using BFS over thoughts.
        
        Args:
            problem: The problem statement
            is_solution: Function to check if a state is a solution
        
        Returns:
            Dict with best solution, path, and all explored nodes
        """
        # Initialize with root node
        root = ThoughtNode(
            thought="Start",
            state=f"Problem: {problem}\n\nLet me think step by step.",
            score=1.0,
            depth=0
        )
        
        current_level = [root]
        all_nodes = [root]
        best_solution = None
        best_score = -float('inf')
        
        for depth in range(self.max_depth):
            print(f"\nDepth {depth + 1}: Exploring {len(current_level)} nodes...")
            
            next_level = []
            
            for node in current_level:
                # Generate candidate thoughts
                thoughts = generate_thoughts(
                    node.state, 
                    n_thoughts=self.n_candidates
                )
                
                for thought in thoughts:
                    # Create new state
                    new_state = f"{node.state}\n\nStep {depth + 1}: {thought}"
                    
                    # Evaluate
                    score = evaluate_thought(node.state, thought)
                    
                    # Create child node
                    child = ThoughtNode(
                        thought=thought,
                        state=new_state,
                        score=score,
                        depth=depth + 1,
                        parent=node
                    )
                    node.children.append(child)
                    all_nodes.append(child)
                    next_level.append(child)
                    
                    # Check if solution
                    if is_solution and is_solution(new_state):
                        if score > best_score:
                            best_solution = child
                            best_score = score
            
            # Keep only top-k nodes for next level
            next_level.sort(key=lambda n: n.score, reverse=True)
            current_level = next_level[:self.beam_width]
            
            print(f"  Generated {len(next_level)} candidates, kept {len(current_level)}")
            for node in current_level:
                print(f"    Score {node.score:.2f}: {node.thought[:50]}...")
        
        # If no explicit solution found, return best final node
        if best_solution is None and current_level:
            best_solution = max(current_level, key=lambda n: n.score)
        
        return {
            "solution": best_solution,
            "path": best_solution.get_path() if best_solution else [],
            "final_state": best_solution.state if best_solution else None,
            "all_nodes": all_nodes,
            "nodes_explored": len(all_nodes)
        }


# Create ToT solver
tot_bfs = TreeOfThoughtsBFS(
    model, tokenizer, device=device,
    n_candidates=2,  # Generate 2 candidates per node
    beam_width=2,    # Keep top 2 at each level
    max_depth=2      # Go 2 levels deep
)

print("BFS Tree of Thoughts solver ready.")

# Test ToT-BFS on a simple problem
problem = "What is 15 + 28? Show your reasoning."

print(f"Problem: {problem}")
print("\n" + "="*60)

result = tot_bfs.solve(problem)

print("\n" + "="*60)
print("SOLUTION PATH:")
print("="*60)
for i, step in enumerate(result["path"]):
    print(f"Step {i}: {step}")

print(f"\nTotal nodes explored: {result['nodes_explored']}")

class TreeOfThoughtsDFS:
    """
    Tree of Thoughts with Depth-First Search.
    
    Explores deeply first, backtracks on failure.
    """
    
    def __init__(self, model, tokenizer, device="cuda",
                 n_candidates: int = 3, max_depth: int = 4,
                 prune_threshold: float = 0.2):
        self.model = model
        self.tokenizer = tokenizer
        self.device = device
        self.n_candidates = n_candidates
        self.max_depth = max_depth
        self.prune_threshold = prune_threshold  # Don't explore if score < this
        self.nodes_explored = 0
    
    def dfs(self, node: ThoughtNode, is_solution: Callable = None) -> Optional[ThoughtNode]:
        """
        Recursive DFS exploration.
        
        Returns the first solution found, or None.
        """
        self.nodes_explored += 1
        
        # Check depth limit
        if node.depth >= self.max_depth:
            return node if is_solution is None or is_solution(node.state) else None
        
        # Check if already a solution
        if is_solution and is_solution(node.state):
            return node
        
        # Generate and evaluate candidates
        thoughts = generate_thoughts(node.state, n_thoughts=self.n_candidates)
        candidates = []
        
        for thought in thoughts:
            new_state = f"{node.state}\n\nStep {node.depth + 1}: {thought}"
            score = evaluate_thought(node.state, thought)
            
            # Prune low-scoring candidates
            if score < self.prune_threshold:
                continue
                
            candidates.append((thought, new_state, score))
        
        # Sort by score, explore best first
        candidates.sort(key=lambda x: x[2], reverse=True)
        
        for thought, new_state, score in candidates:
            child = ThoughtNode(
                thought=thought,
                state=new_state,
                score=score,
                depth=node.depth + 1,
                parent=node
            )
            node.children.append(child)
            
            # Recurse
            result = self.dfs(child, is_solution)
            if result is not None:
                return result
        
        # No solution found in this subtree
        return None
    
    def solve(self, problem: str, is_solution: Callable = None) -> dict:
        """
        Solve a problem using DFS.
        """
        self.nodes_explored = 0
        
        root = ThoughtNode(
            thought="Start",
            state=f"Problem: {problem}\n\nLet me think step by step.",
            score=1.0,
            depth=0
        )
        
        solution = self.dfs(root, is_solution)
        
        return {
            "solution": solution,
            "path": solution.get_path() if solution else [],
            "final_state": solution.state if solution else None,
            "nodes_explored": self.nodes_explored
        }


# Create DFS solver
tot_dfs = TreeOfThoughtsDFS(
    model, tokenizer, device=device,
    n_candidates=2,
    max_depth=3,
    prune_threshold=0.2
)

print("DFS Tree of Thoughts solver ready.")

# Compare BFS vs DFS
problem = "A store has 50 items. They sell 15 and receive 20 more. How many items do they have?"

print(f"Problem: {problem}")
print(f"Correct answer: 50 - 15 + 20 = 55")

print("\n" + "="*60)
print("BFS APPROACH:")
print("="*60)
bfs_result = tot_bfs.solve(problem)
print(f"\nNodes explored: {bfs_result['nodes_explored']}")

print("\n" + "="*60)
print("DFS APPROACH:")
print("="*60)
dfs_result = tot_dfs.solve(problem)
print(f"\nNodes explored: {dfs_result['nodes_explored']}")

def prompt_based_tot(problem: str, n_experts: int = 3) -> str:
    """
    Tree of Thoughts through clever prompting.
    
    No explicit search algorithm — just ask the model to simulate
    multiple reasoning paths with self-critique.
    """
    prompt = f"""Imagine {n_experts} different experts are answering this question.
All experts will write down 1 step of their thinking, then share it with the group.
Then all experts will go on to the next step, etc.
If any expert realizes they're wrong at any point, they leave.
Continue until only one expert remains or all agree on the answer.

Question: {problem}

Expert 1:"""
    
    response = generate_text(prompt, max_new_tokens=300, temperature=0.7)
    return prompt + response


# Test prompt-based ToT
problem = "If I have 3 apples and buy 2 bags with 4 apples each, how many apples do I have?"

print("="*60)
print("PROMPT-BASED TREE OF THOUGHTS")
print("="*60)
print(f"Problem: {problem}")
print(f"Correct: 3 + 2×4 = 11")
print("\n" + "-"*60 + "\n")

result = prompt_based_tot(problem, n_experts=3)
# Just show the response part
print(result.split("Expert 1:")[1] if "Expert 1:" in result else result)