Chain-of-Thought Prompting

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# We'll use Qwen2.5-1.5B-Instruct — small enough to run locally,
# but capable enough to actually demonstrate reasoning
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")

# Move to GPU if available
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
model.eval()

print(f"Loaded on {device}")

def generate_response(prompt, max_new_tokens=100, temperature=0.7):
    """
    Generate a response from the model.
    
    Args:
        prompt: The input text
        max_new_tokens: How many tokens to generate
        temperature: Creativity (0 = deterministic, 1 = creative)
    
    Returns:
        The generated text (excluding the prompt)
    """
    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,
        )
    
    # Decode and remove the prompt from the output
    full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    response = full_response[len(prompt):].strip()
    
    return response

# Test it
test_prompt = "The capital of France is"
print(f"Prompt: {test_prompt}")
print(f"Response: {generate_response(test_prompt, max_new_tokens=20)}")

# A simple arithmetic problem
problem = "If a store has 45 apples and sells 12, then receives a shipment of 30 more, how many apples does it have?"

# Approach 1: Direct prompting (just ask for the answer)
direct_prompt = f"""Question: {problem}
Answer:"""

# Approach 2: Chain-of-Thought (zero-shot)
cot_prompt = f"""Question: {problem}
Answer: Let's think step by step."""

print("="*70)
print("DIRECT PROMPTING")
print("="*70)
print(f"Prompt: {direct_prompt}")
print(f"\nResponse: {generate_response(direct_prompt, max_new_tokens=50)}")

print("\n")
print("="*70)
print("CHAIN-OF-THOUGHT PROMPTING")
print("="*70)
print(f"Prompt: {cot_prompt}")
print(f"\nResponse: {generate_response(cot_prompt, max_new_tokens=100)}")

print("\n")
print("The correct answer is: 45 - 12 + 30 = 63 apples")

# Few-shot examples demonstrating step-by-step reasoning
few_shot_prompt = """Solve these math problems step by step.

Question: A baker has 24 cupcakes. She sells 8 and then bakes 12 more. How many cupcakes does she have?
Answer: Let's work through this step by step.
1. The baker starts with 24 cupcakes.
2. She sells 8 cupcakes: 24 - 8 = 16 cupcakes.
3. She bakes 12 more: 16 + 12 = 28 cupcakes.
Therefore, the baker has 28 cupcakes.

Question: Tom has 15 marbles. He gives 4 to his friend and finds 7 more. How many marbles does Tom have now?
Answer: Let's work through this step by step.
1. Tom starts with 15 marbles.
2. He gives away 4 marbles: 15 - 4 = 11 marbles.
3. He finds 7 more: 11 + 7 = 18 marbles.
Therefore, Tom has 18 marbles.

Question: If a store has 45 apples and sells 12, then receives a shipment of 30 more, how many apples does it have?
Answer: Let's work through this step by step."""

print("="*70)
print("FEW-SHOT CHAIN-OF-THOUGHT")
print("="*70)
response = generate_response(few_shot_prompt, max_new_tokens=100)
print(f"Response: {response}")
print("\nCorrect answer: 63 apples")

class ChainOfThoughtPrompter:
    """
    Wraps a language model with Chain-of-Thought prompting.
    
    Supports both zero-shot and few-shot modes.
    """
    
    def __init__(self, model, tokenizer, device="cuda"):
        self.model = model
        self.tokenizer = tokenizer
        self.device = device
        self.few_shot_examples = []
    
    def add_example(self, question: str, reasoning: str, answer: str):
        """
        Add a few-shot example.
        
        Args:
            question: The problem statement
            reasoning: Step-by-step solution
            answer: Final answer
        """
        self.few_shot_examples.append({
            "question": question,
            "reasoning": reasoning,
            "answer": answer
        })
    
    def build_prompt(self, question: str, zero_shot: bool = False) -> str:
        """
        Build the full prompt with examples (if any).
        
        Args:
            question: The question to answer
            zero_shot: If True, just use "Let's think step by step"
        
        Returns:
            The formatted prompt
        """
        if zero_shot or len(self.few_shot_examples) == 0:
            # Zero-shot CoT
            return f"Question: {question}\nAnswer: Let's think step by step."
        
        # Few-shot CoT
        prompt_parts = []
        
        for ex in self.few_shot_examples:
            prompt_parts.append(
                f"Question: {ex['question']}\n"
                f"Answer: {ex['reasoning']}\n"
                f"Therefore, the answer is: {ex['answer']}"
            )
        
        # Add the new question
        prompt_parts.append(
            f"Question: {question}\n"
            f"Answer:"
        )
        
        return "\n\n".join(prompt_parts)
    
    def solve(self, question: str, zero_shot: bool = False, 
              max_new_tokens: int = 150, temperature: float = 0.7) -> dict:
        """
        Solve a problem using chain-of-thought.
        
        Args:
            question: The problem to solve
            zero_shot: Use zero-shot CoT (ignore examples)
            max_new_tokens: Max tokens to generate
            temperature: Sampling temperature
        
        Returns:
            Dict with 'prompt', 'reasoning', and 'full_response'
        """
        prompt = self.build_prompt(question, zero_shot=zero_shot)
        
        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
        
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                do_sample=True,
                pad_token_id=self.tokenizer.eos_token_id,
            )
        
        full_response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        reasoning = full_response[len(prompt):].strip()
        
        return {
            "prompt": prompt,
            "reasoning": reasoning,
            "full_response": full_response
        }


# Create our CoT prompter
cot = ChainOfThoughtPrompter(model, tokenizer, device=device)

# Add some few-shot examples
cot.add_example(
    question="A store has 50 books. They sell 15 and receive 20 more. How many books do they have?",
    reasoning="Let's solve this step by step. Start: 50 books. Sold 15: 50 - 15 = 35. Received 20: 35 + 20 = 55.",
    answer="55 books"
)

cot.add_example(
    question="Maria has $30. She spends $12 on lunch and $8 on a book. How much money does she have left?",
    reasoning="Let's solve this step by step. Start: $30. Spent on lunch: $30 - $12 = $18. Spent on book: $18 - $8 = $10.",
    answer="$10"
)

print("CoT prompter created with 2 few-shot examples.")

# Test our CoT prompter
test_question = "A farmer has 80 chickens. He sells 25 and then buys 40 more. How many chickens does he have?"

print("="*70)
print("FEW-SHOT CoT RESULT")
print("="*70)

result = cot.solve(test_question, zero_shot=False)
print(f"Question: {test_question}")
print(f"\nGenerated reasoning:\n{result['reasoning']}")
print(f"\nCorrect answer: 80 - 25 + 40 = 95 chickens")

import re

def extract_answer(reasoning: str) -> str:
    """
    Try to extract the final answer from a reasoning chain.

    Looks for common patterns like:
    - "the answer is X"
    - "therefore X"
    - "X apples/books/etc"

    Args:
        reasoning: The chain-of-thought output

    Returns:
        The extracted answer, or None if not found
    """
    # Common answer patterns (order matters - explicit answers first)
    patterns = [
        r"(?:the answer is|answer:)\s*(\d+(?:\.\d+)?)",
        r"(\d+(?:\.\d+)?)\s*(?:apples|books|dollars|chickens|marbles|items)",
        r"(?:therefore|thus|so)[,\s]+(?:.*?has\s+)?(\d+(?:\.\d+)?)",
        r"total[:\s]+(\d+(?:\.\d+)?)",
    ]

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

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

    return None


# Test the extractor
test_chains = [
    "Let's solve this. 50 - 15 = 35. 35 + 20 = 55. The answer is 55.",
    "Start with 80 chickens. Sell 25: 80 - 25 = 55. Buy 40: 55 + 40 = 95 chickens.",
    "Therefore, the farmer has 95 chickens.",
    "Total: 150 items in the store.",
]

print("Testing answer extraction:")
print("="*60)
for chain in test_chains:
    answer = extract_answer(chain)
    print(f"Chain: {chain[:50]}...")
    print(f"Extracted: {answer}")
    print()