Self-Consistency

import torch
import numpy as np
from transformers import AutoTokenizer, AutoModelForCausalLM
from collections import Counter
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}")

def extract_number(text: str) -> str:
    """
    Extract the final numerical answer from reasoning text.

    Tries several patterns, falls back to last number found.
    """
    # Try common answer patterns (order matters - more specific first)
    patterns = [
        r"(?:the answer is|answer:)\s*(\d+(?:\.\d+)?)",
        r"(?:therefore|thus|so)[,\s]+(?:the answer is\s*)?(\d+(?:\.\d+)?)",
        r"=\s*(\d+(?:\.\d+)?)\s*(?:\.|$)",
        r"(\d+(?:\.\d+)?)\s*(?:total|in all|altogether)",
    ]

    text_lower = text.lower()
    for pattern in patterns:
        match = re.search(pattern, text_lower)
        if match:
            return match.group(1)

    # Fallback: last number in the text
    numbers = re.findall(r"\b\d+\.?\d*\b", text)
    if numbers:
        return numbers[-1]

    return None


def generate_chain(prompt: str, max_new_tokens: int = 100,
                   temperature: float = 0.7) -> str:
    """
    Generate a single reasoning chain.

    Args:
        prompt: The CoT prompt
        max_new_tokens: Max tokens to generate
        temperature: Sampling temperature (>0 for diversity)

    Returns:
        The generated reasoning chain
    """
    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()


# Test single generation
test_prompt = "Question: What is 5 + 7?\nAnswer: Let's think step by step."
chain = generate_chain(test_prompt)
print(f"Generated chain: {chain}")
print(f"Extracted answer: {extract_number(chain)}")

class SelfConsistency:
    """
    Self-Consistency for chain-of-thought reasoning.
    
    Sample multiple reasoning paths, extract answers, take majority vote.
    """
    
    def __init__(self, model, tokenizer, device="cuda"):
        self.model = model
        self.tokenizer = tokenizer
        self.device = device
    
    def sample_chains(self, prompt: str, n_samples: int = 5,
                      max_new_tokens: int = 100,
                      temperature: float = 0.7) -> list:
        """
        Generate N independent reasoning chains.
        
        Args:
            prompt: The CoT prompt
            n_samples: Number of chains to generate
            max_new_tokens: Max tokens per chain
            temperature: Sampling temperature (higher = more diversity)
        
        Returns:
            List of (chain, extracted_answer) tuples
        """
        chains = []
        
        for i in range(n_samples):
            chain = generate_chain(
                prompt, 
                max_new_tokens=max_new_tokens,
                temperature=temperature
            )
            answer = extract_number(chain)
            chains.append((chain, answer))
        
        return chains
    
    def majority_vote(self, answers: list) -> tuple:
        """
        Take majority vote across answers.
        
        Args:
            answers: List of extracted answers (some may be None)
        
        Returns:
            (most_common_answer, vote_count, total_valid)
        """
        # Filter out None answers
        valid_answers = [a for a in answers if a is not None]
        
        if not valid_answers:
            return None, 0, 0
        
        # Count votes
        counter = Counter(valid_answers)
        most_common, count = counter.most_common(1)[0]
        
        return most_common, count, len(valid_answers)
    
    def solve(self, question: str, n_samples: int = 5,
              max_new_tokens: int = 100,
              temperature: float = 0.7) -> dict:
        """
        Solve a problem using self-consistency.
        
        Args:
            question: The problem to solve
            n_samples: Number of reasoning chains to sample
            max_new_tokens: Max tokens per chain
            temperature: Sampling temperature
        
        Returns:
            Dict with answer, confidence, chains, and vote details
        """
        # Build CoT prompt
        prompt = f"Question: {question}\nAnswer: Let's think step by step."
        
        # Sample chains
        chains = self.sample_chains(
            prompt, 
            n_samples=n_samples,
            max_new_tokens=max_new_tokens,
            temperature=temperature
        )
        
        # Extract answers and vote
        answers = [answer for _, answer in chains]
        final_answer, votes, total_valid = self.majority_vote(answers)
        
        # Calculate confidence as vote fraction
        confidence = votes / total_valid if total_valid > 0 else 0.0
        
        return {
            "answer": final_answer,
            "confidence": confidence,
            "votes": votes,
            "total_samples": n_samples,
            "valid_samples": total_valid,
            "chains": chains,
            "vote_distribution": Counter(answers)
        }


# Create self-consistency solver
sc = SelfConsistency(model, tokenizer, device=device)
print("Self-consistency solver ready.")

# Test self-consistency on a math problem
# Note: With a small 1.5B model, the model often gets multi-step arithmetic wrong.
# This actually demonstrates WHY self-consistency matters - and its limitations
# when the base model accuracy is below 50%!

question = "A store has 45 books. They sell 12 and receive 30 more. How many books do they have?"

print(f"Question: {question}")
print(f"Correct answer: 45 - 12 + 30 = 63")
print("\n" + "="*70)
print("Sampling 5 reasoning chains...")
print("="*70 + "\n")

result = sc.solve(question, n_samples=5, temperature=0.8)

# Show each chain
for i, (chain, answer) in enumerate(result["chains"], 1):
    print(f"Chain {i}: {chain[:80]}...")
    print(f"  → Extracted answer: {answer}")
    print()

print("="*70)
print(f"Vote distribution: {dict(result['vote_distribution'])}")
print(f"Final answer: {result['answer']}")
print(f"Confidence: {result['confidence']:.1%} ({result['votes']}/{result['valid_samples']} votes)")

from scipy import stats
import matplotlib.pyplot as plt

def majority_vote_accuracy(p, n):
    """
    Calculate probability that majority vote gives correct answer.
    
    Args:
        p: Probability each individual sample is correct
        n: Number of samples
    
    Returns:
        Probability of correct majority
    """
    # Need more than half to be correct
    # For n samples, we need at least ceil(n/2) correct
    k_min = (n // 2) + 1
    
    # Sum probability of getting k_min, k_min+1, ..., n correct
    return 1 - stats.binom.cdf(k_min - 1, n, p)


# Plot the improvement from self-consistency
sample_sizes = [1, 3, 5, 9, 15, 21, 41]
base_accuracies = [0.4, 0.5, 0.6, 0.7, 0.8]

plt.figure(figsize=(10, 6))

for p in base_accuracies:
    accuracies = [majority_vote_accuracy(p, n) for n in sample_sizes]
    plt.plot(sample_sizes, accuracies, 'o-', label=f'Base p={p}', linewidth=2, markersize=8)

plt.axhline(y=0.5, color='gray', linestyle='--', alpha=0.5)
plt.xlabel('Number of Samples (N)', fontsize=12)
plt.ylabel('Probability of Correct Majority', fontsize=12)
plt.title('Self-Consistency Improves Accuracy\n(as long as base accuracy > 50%)', fontsize=14)
plt.legend(title='Single-sample accuracy')
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

print("Key insight: If the model is right more often than not (p > 0.5),")
print("sampling more chains and voting will push accuracy toward 100%.")
print("\nBut if p < 0.5, voting makes things WORSE!")

# Demonstrate the effect of temperature on diversity
question = "What is 15 + 27?"
prompt = f"Question: {question}\nAnswer: Let's think step by step."

temperatures = [0.3, 0.7, 1.0]

for temp in temperatures:
    print(f"\n{'='*60}")
    print(f"Temperature = {temp}")
    print("="*60)
    
    chains = []
    for i in range(3):
        chain = generate_chain(prompt, max_new_tokens=50, temperature=temp)
        chains.append(chain)
        print(f"\nSample {i+1}: {chain[:60]}...")
    
    # Check diversity: how many unique answers?
    answers = [extract_number(c) for c in chains]
    unique = len(set(a for a in answers if a is not None))
    print(f"\nUnique answers: {unique}/3")

def generate_chain_with_score(prompt: str, max_new_tokens: int = 100,
                               temperature: float = 0.7) -> tuple:
    """
    Generate a chain and compute its log probability.
    
    The log probability indicates how "confident" the model was
    when generating this particular chain.
    
    Returns:
        (chain_text, log_probability)
    """
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    input_len = inputs["input_ids"].shape[1]
    
    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,
            return_dict_in_generate=True,
            output_scores=True,
        )
    
    # Get generated token IDs (excluding prompt)
    generated_ids = outputs.sequences[0, input_len:]
    
    # Compute log probability of the generated sequence
    # scores[i] has shape (batch, vocab) for token i
    log_prob = 0.0
    for i, token_id in enumerate(generated_ids):
        if i >= len(outputs.scores):
            break
        logits = outputs.scores[i][0]  # Shape: (vocab_size,)
        log_probs = torch.log_softmax(logits, dim=-1)
        log_prob += log_probs[token_id].item()
    
    # Normalize by length (average log prob per token)
    avg_log_prob = log_prob / max(len(generated_ids), 1)
    
    # Decode the chain
    chain = tokenizer.decode(generated_ids, skip_special_tokens=True)
    
    return chain, avg_log_prob


def weighted_majority_vote(chains_with_scores: list) -> tuple:
    """
    Weighted majority vote based on log probabilities.
    
    Args:
        chains_with_scores: List of (chain, answer, log_prob) tuples
    
    Returns:
        (best_answer, weight_sum, total_weight)
    """
    from collections import defaultdict
    
    # Convert log probs to weights (using softmax for stability)
    log_probs = [lp for _, _, lp in chains_with_scores]
    
    # Softmax: exp(log_prob) / sum(exp(log_probs))
    max_lp = max(log_probs)
    exp_lps = [np.exp(lp - max_lp) for lp in log_probs]  # Subtract max for stability
    total = sum(exp_lps)
    weights = [e / total for e in exp_lps]
    
    # Weighted vote
    answer_weights = defaultdict(float)
    for (_, answer, _), weight in zip(chains_with_scores, weights):
        if answer is not None:
            answer_weights[answer] += weight
    
    if not answer_weights:
        return None, 0, 0
    
    best_answer = max(answer_weights, key=answer_weights.get)
    return best_answer, answer_weights[best_answer], sum(answer_weights.values())


# Demo weighted voting
print("Generating chains with confidence scores...\n")

question = "What is 8 + 15?"
prompt = f"Question: {question}\nAnswer: Let's think step by step."

chains_with_scores = []
for i in range(5):
    chain, log_prob = generate_chain_with_score(prompt, max_new_tokens=50)
    answer = extract_number(chain)
    chains_with_scores.append((chain, answer, log_prob))
    print(f"Chain {i+1}:")
    print(f"  Text: {chain[:50]}...")
    print(f"  Answer: {answer}")
    print(f"  Avg log prob: {log_prob:.3f}")
    print()

# Compare voting methods
answers = [a for _, a, _ in chains_with_scores]
unweighted_answer, votes, _ = sc.majority_vote(answers)
weighted_answer, weight, _ = weighted_majority_vote(chains_with_scores)

print("="*50)
print(f"Unweighted vote: {unweighted_answer}")
print(f"Weighted vote:   {weighted_answer}")
print(f"Correct answer:  23")

def self_consistency_early_stop(question: str, max_samples: int = 20,
                                 confidence_threshold: float = 0.8,
                                 min_samples: int = 3) -> dict:
    """
    Self-consistency with early stopping.

    Stop sampling when we're confident enough in the answer.

    Args:
        question: The problem to solve
        max_samples: Maximum chains to generate
        confidence_threshold: Stop if vote share exceeds this
        min_samples: Always generate at least this many

    Returns:
        Dict with answer, samples_used, and chains
    """
    prompt = f"Question: {question}\nAnswer: Let's think step by step."

    chains = []
    answers = []

    for i in range(max_samples):
        # Generate one chain
        chain = generate_chain(prompt, max_new_tokens=80, temperature=0.7)
        answer = extract_number(chain)

        chains.append(chain)
        answers.append(answer)

        # Check if we can stop early
        if i + 1 >= min_samples:
            valid = [a for a in answers if a is not None]
            if valid:
                counter = Counter(valid)
                top_answer, top_count = counter.most_common(1)[0]
                confidence = top_count / len(valid)

                if confidence >= confidence_threshold:
                    return {
                        "answer": top_answer,
                        "confidence": confidence,
                        "samples_used": i + 1,
                        "max_samples": max_samples,
                        "early_stopped": True,
                        "chains": list(zip(chains, answers))
                    }

    # Didn't stop early—use all samples
    valid = [a for a in answers if a is not None]
    if valid:
        counter = Counter(valid)
        top_answer, top_count = counter.most_common(1)[0]
        confidence = top_count / len(valid)
    else:
        top_answer, confidence = None, 0.0

    return {
        "answer": top_answer,
        "confidence": confidence,
        "samples_used": max_samples,
        "max_samples": max_samples,
        "early_stopped": False,
        "chains": list(zip(chains, answers))
    }


# Demo early stopping
# Note: The small model often gets this wrong - demonstrating that self-consistency
# only helps when base accuracy > 50%. This motivates the techniques in later notebooks!
print("Testing self-consistency with early stopping...\n")

question = "A farmer has 30 cows. He sells 10 and buys 15 more. How many cows does he have?"
print(f"Question: {question}")
print(f"Correct: 30 - 10 + 15 = 35\n")

result = self_consistency_early_stop(
    question,
    max_samples=10,
    confidence_threshold=0.7
)

print(f"Answer: {result['answer']}")
print(f"Confidence: {result['confidence']:.1%}")
print(f"Samples used: {result['samples_used']}/{result['max_samples']}")
print(f"Early stopped: {result['early_stopped']}")

if result['early_stopped']:
    savings = (result['max_samples'] - result['samples_used']) / result['max_samples']
    print(f"Compute saved: {savings:.0%}")