main.py

from decoder import SpeculativeDecoder
import torch
import time


def test_speculative_decoding():
    """
    Test the speculative decoding implementation and compare with baseline.
    """
    print("Testing Speculative Decoding Implementation")
    print("=" * 60)

    # Initialize decoder with small models for testing
    decoder = SpeculativeDecoder(
        target_model_name="gpt2-large",
        draft_model_name="distilgpt2",
        device="cuda" if torch.cuda.is_available() else "cpu",
        temperature=0.8,
        top_p=0.9,
    )

    # Test prompts
    prompts = ["Once upon a time in a land far away,"]

    for i, prompt in enumerate(prompts):
        print(f"\n{'='*60}")
        print(f"TEST {i+1}: {prompt}")
        print("=" * 60)

        # First, run baseline (target model only)
        print("\n" + "BASELINE: Target Model Only")
        print("=" * 60)
        baseline_text, baseline_time = decoder.generate_baseline(
            prompt=prompt,
            max_length=100,
        )
        print(f"\nBaseline generated text:\n{baseline_text}\n")

        # Then run speculative decoding
        print("\n" + "SPECULATIVE DECODING: Draft + Target")
        print("=" * 60)
        spec_start = time.time()
        spec_text = decoder.generate(
            prompt=prompt,
            max_length=100,
            k=4,
        )
        spec_time = time.time() - spec_start
        print(f"\nSpeculative generated text:\n{spec_text}\n")

        # Compare results
        print("\n" + "COMPARISON")
        print("=" * 60)
        speedup = baseline_time / spec_time if spec_time > 0 else 0
        print(f"Baseline time: {baseline_time:.3f}s")
        print(f"Speculative time: {spec_time:.3f}s")
        print(f"Speedup: {speedup:.2f}x")
        print(
            f"Time saved: {baseline_time - spec_time:.3f}s ({(1 - spec_time/baseline_time)*100:.1f}%)"
        )
        print("=" * 60)


if __name__ == "__main__":
    test_speculative_decoding()