import json
import os
import argparse
from vllm import LLM, SamplingParams

prompt = "Explain the theory of relativity in simple terms."

RESULT_FILE = "eplb_test_output.json"

sampling_params = SamplingParams(
    temperature=0.0,
    top_p=1.0,
    top_k=1,
    max_tokens=4096
)

def run_inference(model_path: str, enable_eplb: bool, num_redundant_experts: int = 0):
    print(f"Running inference with EPLB={enable_eplb}, redundant experts={num_redundant_experts}")
    
    llm = LLM(
        model=model_path,
        tensor_parallel_size=2,
        enable_expert_parallel=True,
        enable_eplb=enable_eplb,
        num_redundant_experts=num_redundant_experts if enable_eplb else 0,
        eplb_window_size=1000,
        eplb_step_interval=100,
        enforce_eager=True,
        trust_remote_code=True
    )
    
    result = llm.generate([prompt], sampling_params)
    output_text = result[0].outputs[0].text.strip()
    
    print("Output:")
    print(output_text)
    print("-" * 50)

    return output_text

def save_result(key: str, value: list):
    if os.path.exists(RESULT_FILE):
        with open(RESULT_FILE, "r") as f:
            results = json.load(f)
    else:
        results = {}

    results[key] = value

    with open(RESULT_FILE, "w") as f:
        json.dump(results, f, indent=2)

    print(f"Output saved to {RESULT_FILE}")

def load_results():
    if os.path.exists(RESULT_FILE):
        with open(RESULT_FILE, "r") as f:
            return json.load(f)
    return {}

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--mode", type=str, choices=["eplb", "normal", "compare"], required=True)
    args = parser.parse_args()

    MODEL_PATH = "baidu/ERNIE-4.5-21B-A3B-PT"

    if args.mode == "eplb":
        outputs = run_inference(MODEL_PATH, enable_eplb=True, num_redundant_experts=32)
        save_result("eplb", outputs)
    elif args.mode == "normal":
        outputs = run_inference(MODEL_PATH, enable_eplb=False)
        save_result("normal", outputs)

{
  "eplb": "The **theory of relativity**, developed by Albert Einstein in the early 20th century, has two main parts: **special relativity** and **general relativity**. Here's a simple breakdown:\n\n### **1. Special Relativity (1905)**\n- **Key Idea**: Time and space are not absolute\u2014they depend on how you're moving.\n- **Key Concepts**:\n  - **Time Dilation**: Moving clocks run slower. For example, if you're on a fast train, your watch might tick slightly slower than someone standing still.\n  - **Length Contraction**: Objects appear shorter when moving. A ruler on a train would look shorter to someone watching from the platform.\n  - **Mass-Energy Equivalence**: \"E=mc\u00b2\" means mass and energy are the same thing. A tiny amount of mass can release a huge amount of energy (like in nuclear reactions).\n- **Why It Matters**: It explains why the speed of light (300,000 km/s) is the same for everyone, no matter how fast they're moving.\n\n### **2. General Relativity (1915)**\n- **Key Idea**: Gravity isn't a force\u2014it's the curvature of spacetime caused by mass and energy.\n- **Key Concepts**:\n  - **Spacetime Curvature**: Massive objects like planets and stars warp the fabric of space and time. This curvature tells objects how to move.\n  - **Gravity as Geometry**: Instead of being pulled by a force, objects follow straight paths (geodesics) in curved spacetime. For example, Earth orbits the Sun because spacetime is curved by the Sun's mass.\n- **Why It Matters**: It explains why light bends around massive objects (like stars during eclipses) and predicts black holes.\n\n### **Simple Analogy**\nImagine spacetime as a trampoline:\n- **Special Relativity**: If you jump on the trampoline (move fast), the surface curves more, and things seem to move differently.\n- **General Relativity**: A heavy ball (massive object) on the trampoline creates a deep dent. Smaller balls (objects) roll toward it, just like gravity pulls objects toward massive bodies.\n\n### **Why It Matters Today**\n- **GPS Satellites**: Must account for time dilation (special relativity) because they move fast relative to Earth.\n- **Black Holes & Cosmology**: General relativity predicts phenomena like black holes and the Big Bang.\n- **Quantum Mechanics**: Relativity and quantum physics are still being unified, but they've already revolutionized our understanding of the universe.\n\nIn short, Einstein showed that time, space, and gravity are interconnected, and our perception of reality depends on how we're moving or where we are in the universe.",
  "normal": "The **theory of relativity**, developed by Albert Einstein in the early 20th century, has two main parts: **special relativity** and **general relativity**. Here's a simple breakdown:\n\n### **1. Special Relativity (1905)**\n- **Key Idea**: Time and space are not absolute\u2014they depend on how you're moving.\n- **Key Concepts**:\n  - **Time Dilation**: Moving clocks run slower. For example, if you're on a fast train, your watch might tick slightly slower than someone standing still.\n  - **Length Contraction**: Objects appear shorter when moving. A ruler on a train would look shorter to someone watching from the platform.\n  - **Mass-Energy Equivalence**: \"E=mc\u00b2\" means mass and energy are the same thing. A tiny amount of mass can release a huge amount of energy (like in nuclear reactions).\n- **Why It Matters**: It explains why the speed of light (300,000 km/s) is the same for everyone, no matter how fast they're moving.\n\n### **2. General Relativity (1915)**\n- **Key Idea**: Gravity isn't a force\u2014it's the curvature of spacetime caused by mass and energy.\n- **Key Concepts**:\n  - **Spacetime Curvature**: Massive objects like planets and stars warp the fabric of space and time. This curvature tells objects how to move.\n  - **Gravity as Geometry**: Instead of being pulled by a force, objects follow straight paths (geodesics) in curved spacetime. For example, Earth orbits the Sun because spacetime is curved by the Sun's mass.\n- **Why It Matters**: It explains why light bends around massive objects (like stars during eclipses) and predicts black holes.\n\n### **Simple Analogy**\nImagine spacetime as a trampoline:\n- **Special Relativity**: If you jump on the trampoline (move fast), the fabric curves more, and nearby objects (like a ball) move along the curves.\n- **General Relativity**: A heavy ball (like the Sun) sinks deeply into the trampoline, curving the fabric so that a smaller ball (like Earth) rolls around it in a straight line (but on the curved surface).\n\n### **Why It's Revolutionary**\n- It unified space and time (special relativity) and showed gravity is geometry (general relativity).\n- It predicted phenomena like black holes, gravitational waves, and the expansion of the universe.\n\nIn short, relativity reshaped our understanding of reality, showing that time, space, and gravity are interconnected in ways that defy everyday intuition. \ud83c\udf0c\ud83d\ude80"
}