app.py

"""
Weak Supervision Reasoning Explorer
Based on: "When Can LLMs Learn to Reason with Weak Supervision?" (arXiv:2604.18574)
"""

import gradio as gr
import random

# Sample reasoning problems for demonstration
SAMPLE_PROBLEMS = [
    {
        "id": 1,
        "question": "A farmer has 17 sheep and all but 9 die. How many are left?",
        "correct_answer": "9",
        "requires_reasoning": True,
        "category": "Word Problem"
    },
    {
        "id": 2,
        "question": "What is the sum of numbers from 1 to 10?",
        "correct_answer": "55",
        "requires_reasoning": True,
        "category": "Mathematics"
    },
    {
        "id": 3,
        "question": "If it takes 5 machines 5 minutes to make 5 widgets, how long does it take 100 machines to make 100 widgets?",
        "correct_answer": "5 minutes",
        "requires_reasoning": True,
        "category": "Logic"
    },
    {
        "id": 4,
        "question": "A bat and a ball cost $11 in total. The bat costs $10 more than the ball. How much does the ball cost?",
        "correct_answer": "$0.50",
        "requires_reasoning": True,
        "category": "Cognitive Reflection"
    },
    {
        "id": 5,
        "question": "In a lake, there is a patch of lily pads. Every day, the patch doubles in size. If it takes 48 days to cover the entire lake, how long would it take to cover half the lake?",
        "correct_answer": "47 days",
        "requires_reasoning": True,
        "category": "Exponential Growth"
    }
]


def generate_saturation_curve(model_type, weak_supervision_level):
    """Generate simulated reward saturation curves based on paper findings"""
    
    # Training steps
    steps = list(range(0, 1001, 50))
    
    if model_type == "Generalizing Model":
        # Prolonged pre-saturation: reward and performance climb together
        train_reward = [min(0.3 + (s/1000) * 0.6 + random.uniform(-0.05, 0.05), 0.95) for s in steps]
        downstream_perf = [min(0.25 + (s/1000) * 0.65 + random.uniform(-0.03, 0.03), 0.90) for s in steps]
        saturation_point = "Step 800+ (gradual)"
    else:  # Memorizing Model
        # Rapid saturation: reward saturates early, performance plateaus
        train_reward = [min(0.3 + (s/300) * 0.7, 0.98) for s in steps]
        train_reward = [r if i < 6 else 0.95 + random.uniform(-0.02, 0.02) for i, r in enumerate(train_reward)]
        downstream_perf = [min(0.25 + (s/300) * 0.4, 0.55) for s in steps]
        downstream_perf = [p if i < 6 else 0.50 + random.uniform(-0.03, 0.03) for i, p in enumerate(downstream_perf)]
        saturation_point = "Step 300 (rapid)"
    
    # Apply weak supervision degradation
    if weak_supervision_level == "Scarce Data (10%)":
        downstream_perf = [p * 0.7 for p in downstream_perf]
    elif weak_supervision_level == "Noisy Rewards (30% noise)":
        downstream_perf = [p * 0.8 + random.uniform(-0.05, 0.05) for p in downstream_perf]
    elif weak_supervision_level == "Self-Supervised Proxy":
        downstream_perf = [p * 0.75 for p in downstream_perf]
    
    return steps, train_reward, downstream_perf, saturation_point


def analyze_reasoning_faithfulness(answer, reasoning_steps):
    """Analyze if intermediate steps logically support final answer"""
    
    # Simulated faithfulness scoring
    if not reasoning_steps:
        return 0.0, "No reasoning provided"
    
    # Check for key indicators of faithfulness
    indicators = [
        "step" in reasoning_steps.lower(),
        "because" in reasoning_steps.lower() or "since" in reasoning_steps.lower(),
        "therefore" in reasoning_steps.lower() or "so" in reasoning_steps.lower(),
        len(reasoning_steps.split()) > 20  # Substantial explanation
    ]
    
    faithfulness_score = sum(indicators) / len(indicators)
    
    if faithfulness_score >= 0.75:
        assessment = "High faithfulness: Steps logically support conclusion"
    elif faithfulness_score >= 0.5:
        assessment = "Moderate faithfulness: Some logical gaps"
    else:
        assessment = "Low faithfulness: Steps don't clearly support answer"
    
    return faithfulness_score, assessment


def compare_weak_supervision_scenarios():
    """Compare different weak supervision settings"""
    
    scenarios = [
        {
            "name": "Scarce Data (10%)",
            "description": "Only 10% of training data available",
            "generalizing_perf": "72%",
            "memorizing_perf": "35%",
            "key_factor": "SFT on reasoning traces critical"
        },
        {
            "name": "Noisy Rewards (30% noise)",
            "description": "30% of reward signals are incorrect",
            "generalizing_perf": "68%",
            "memorizing_perf": "42%",
            "key_factor": "Faithfulness filters noise"
        },
        {
            "name": "Self-Supervised Proxy",
            "description": "No external rewards, self-generated signals",
            "generalizing_perf": "65%",
            "memorizing_perf": "28%",
            "key_factor": "Internal consistency matters"
        }
    ]
    
    output = "## Weak Supervision Scenarios Comparison\n\n"
    for s in scenarios:
        output += f"### {s['name']}\n"
        output += f"**Setting:** {s['description']}\n\n"
        output += f"| Model Type | Performance |\n"
        output += f"|------------|-------------|\n"
        output += f"| Generalizing | {s['generalizing_perf']} |\n"
        output += f"| Memorizing | {s['memorizing_perf']} |\n\n"
        output += f"**Key Factor:** {s['key_factor']}\n\n---\n\n"
    
    return output


def demonstrate_saturation_dynamics(model_type, weak_supervision):
    """Demonstrate reward saturation dynamics"""
    
    steps, train_reward, downstream_perf, saturation = generate_saturation_curve(model_type, weak_supervision)
    
    output = f"## {model_type} - {weak_supervision}\n\n"
    output += f"**Saturation Point:** {saturation}\n\n"
    output += "### Training Dynamics\n\n"
    output += "| Step | Training Reward | Downstream Performance | Gap |\n"
    output += "|------|-----------------|------------------------|-----|\n"
    
    for i in range(0, len(steps), 2):  # Sample every other step for readability
        gap = train_reward[i] - downstream_perf[i]
        output += f"| {steps[i]} | {train_reward[i]:.3f} | {downstream_perf[i]:.3f} | {gap:.3f} |\n"
    
    # Analysis
    if model_type == "Generalizing Model":
        output += "\n### Analysis\n"
        output += "✅ **Prolonged pre-saturation**: Training reward and downstream performance climb together\n"
        output += "✅ **Low gap**: Model is actually learning, not memorizing\n"
        output += "✅ **Generalization**: Performance transfers to downstream tasks\n"
    else:
        output += "\n### Analysis\n"
        output += "❌ **Rapid saturation**: Training reward saturates early\n"
        output += "❌ **High gap**: Large divergence indicates memorization\n"
        output += "❌ **Poor generalization**: Downstream performance plateaus\n"
    
    return output


def evaluate_reasoning_faithfulness(question_id, answer, reasoning):
    """Evaluate reasoning faithfulness for a specific problem"""
    
    problem = SAMPLE_PROBLEMS[int(question_id) - 1]
    
    score, assessment = analyze_reasoning_faithfulness(answer, reasoning)
    
    output = f"## Reasoning Faithfulness Analysis\n\n"
    output += f"**Question:** {problem['question']}\n\n"
    output += f"**Your Answer:** {answer}\n"
    output += f"**Correct Answer:** {problem['correct_answer']}\n\n"
    output += f"**Faithfulness Score:** {score:.2f}/1.0\n\n"
    output += f"**Assessment:** {assessment}\n\n"
    
    output += "### Key Insights from Paper\n\n"
    output += "> 'Reasoning faithfulness, defined as the extent to which intermediate steps logically support the final answer, is the pre-RL property that predicts which regime a model falls into.'\n\n"
    
    if score >= 0.75:
        output += "✅ **High faithfulness** correlates with generalization under weak supervision\n"
    elif score >= 0.5:
        output += "⚠️ **Moderate faithfulness** - may struggle with scarce data or noisy rewards\n"
    else:
        output += "❌ **Low faithfulness** - likely to memorize rather than learn\n"
    
    return output


# Gradio Interface
demo = gr.Blocks(title="Weak Supervision Reasoning Explorer")

with demo:
    gr.Markdown("""
    # 🔬 Weak Supervision Reasoning Explorer
    
    Interactive exploration of **"When Can LLMs Learn to Reason with Weak Supervision?"** (arXiv:2604.18574)
    
    **Core Finding:** Models that generalize exhibit prolonged pre-saturation phases, while rapid saturation indicates memorization.
    """)
    
    with gr.Tab("Reward Saturation Dynamics"):
        with gr.Row():
            with gr.Column():
                model_type = gr.Dropdown(
                    choices=["Generalizing Model", "Memorizing Model"],
                    value="Generalizing Model",
                    label="Model Type"
                )
                weak_supervision = gr.Dropdown(
                    choices=["Full Supervision", "Scarce Data (10%)", "Noisy Rewards (30% noise)", "Self-Supervised Proxy"],
                    value="Scarce Data (10%)",
                    label="Weak Supervision Setting"
                )
                run_btn = gr.Button("Analyze Saturation", variant="primary")
            
            with gr.Column():
                saturation_output = gr.Markdown()
        
        run_btn.click(
            fn=demonstrate_saturation_dynamics,
            inputs=[model_type, weak_supervision],
            outputs=[saturation_output]
        )
    
    with gr.Tab("Weak Supervision Scenarios"):
        gr.Markdown("Compare how different weak supervision settings affect model performance")
        compare_btn = gr.Button("Compare Scenarios", variant="primary")
        comparison_output = gr.Markdown()
        compare_btn.click(fn=compare_weak_supervision_scenarios, outputs=[comparison_output])
    
    with gr.Tab("Reasoning Faithfulness"):
        with gr.Row():
            with gr.Column():
                problem_select = gr.Dropdown(
                    choices=[(f"Problem {p['id']}: {p['category']}", str(p['id'])) for p in SAMPLE_PROBLEMS],
                    value="1",
                    label="Select Problem"
                )
                user_answer = gr.Textbox(label="Your Answer", placeholder="Enter your answer...")
                user_reasoning = gr.TextArea(
                    label="Your Reasoning Steps",
                    placeholder="Explain your reasoning step by step...",
                    lines=5
                )
                eval_btn = gr.Button("Evaluate Faithfulness", variant="primary")
            
            with gr.Column():
                faithfulness_output = gr.Markdown()
        
        eval_btn.click(
            fn=evaluate_reasoning_faithfulness,
            inputs=[problem_select, user_answer, user_reasoning],
            outputs=[faithfulness_output]
        )
    
    gr.Markdown("""
    ---
    
    ### 📄 Paper Reference
    
    **When Can LLMs Learn to Reason with Weak Supervision?**  
    Salman Rahman, Jingyan Shen, Anna Mordvina, Hamid Palangi, Saadia Gabriel, Pavel Izmailov  
    *NYU, Microsoft Research*  
    arXiv:2604.18574
    """)

if __name__ == "__main__":
    demo.launch()