iclr_exp_llama3_eagle1_govreport.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os, json
os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # match your original script

import numpy as np
from tqdm import tqdm
from termcolor import colored
from accelerate import Accelerator
from eagle.model_eagle_l3 import EaModel_L3 as EaModel
import torch
from pathlib import Path

def build_messages_for_task(text: str):
    # keep it simple and neutral
    return [{"role": "user", "content": text}]

def make_input_ids(tokenizer, text: str, device, *, use_chat_template=True, add_generation_prompt=True):
    # Prefer the model's tokenizer chat template if present
    if use_chat_template and hasattr(tokenizer, "apply_chat_template") and getattr(tokenizer, "chat_template", None):
        messages = build_messages_for_task(text)
        try:
            return tokenizer.apply_chat_template(
                messages,
                add_generation_prompt=add_generation_prompt,
                return_tensors="pt",
            ).to(device)
        except Exception as e:
            print(f"[warn] apply_chat_template failed ({e}); falling back to encode().")

    # Fallback for older/plain models
    return tokenizer.encode(text, return_tensors="pt", add_special_tokens=True).to(device)


# ─── FIXED CONFIGS (kept as in your original) ──────────────────────────────────
DATASET_NAME = "govreport"
DATA_FOLDER  = os.path.join("data", DATASET_NAME)
OUT_DIR      = "results/ICLR/EAGLE1"

BASE_MODEL   = "meta-llama/Llama-3.1-8B-Instruct"
DRAFT_MODEL  = "yuhuili/EAGLE-LLaMA3.1-Instruct-8B"  # NOTE: not EAGLE3, kept as-is

RETRIEVAL_CHUNK_SIZE   = 32
RETRIEVE_TOP_K         = 64
RETRIEVE_EVERY_N_STEPS = 4

N_PER_FILE = 20   # lines per file to load (same as original)
MAX_NEW_TOKENS_RUN = 256  # per real run
WARMUP_REPEATS = 3
WARMUP_TOKENS  = 5
REPLICATES_PER_RECORD = 1

# ─── LOAD DATA ─────────────────────────────────────────────────────────────────
file_to_data = {}
for fname in os.listdir(DATA_FOLDER):
    full_path = os.path.join(DATA_FOLDER, fname)
    if not os.path.isfile(full_path):
        continue
    records = []
    with open(full_path, "r", encoding="utf-8", errors="replace") as f:
        for i, line in enumerate(f):
            if i >= N_PER_FILE:
                break
            line = line.strip().lstrip("\ufeff")
            if not line:
                continue
            try:
                rec = json.loads(line)
            except json.JSONDecodeError as e:
                print(f"JSON decode error in {fname} line {i}: {e}")
                continue
            records.append(rec)
    if records:
        # your original used "trunc_length" for govreport
        length = records[-1]["trunc_length"]
        file_to_data[length] = records

sorted_lengths = sorted(file_to_data.keys())

# ─── EXPERIMENT RUNNER ─────────────────────────────────────────────────────────
def run_experiment(use_specextend: bool):
    # Build exp name exactly like your original
    def after_first_slash(x: str) -> str:
        return x.split("/", 1)[-1]
    dataset      = after_first_slash(DATASET_NAME)
    target_model = after_first_slash(BASE_MODEL)
    draft_model  = after_first_slash(DRAFT_MODEL)
    exp_name_base = f"{dataset}_{target_model}_{draft_model}_SpecExtend_{use_specextend}"

    # Ensure unique filename (same logic you had)
    Path(OUT_DIR).mkdir(parents=True, exist_ok=True)
    i = 2
    exp_name = exp_name_base
    out_path = f"{OUT_DIR}/{exp_name}.json"
    while os.path.exists(out_path):
        exp_name = f"{exp_name_base}_{i}"
        out_path = f"{OUT_DIR}/{exp_name}.json"
        i += 1

    exp_config = {
        "exp_name":       exp_name_base,
        "dataset":        dataset,
        "target_model":   target_model,
        "draft_model":    draft_model,
        "use_specextend": use_specextend,
        "cache_parameters": {
            "retrieval_chunk_size":   RETRIEVAL_CHUNK_SIZE,
            "retrieve_top_k":         RETRIEVE_TOP_K,
            "retrieve_every_n_steps": RETRIEVE_EVERY_N_STEPS,
        } if use_specextend else None,
    }

    summary = {}

    # Load model ONCE per experiment (like your conversion request)
    print(colored(f"Loading model (SpecExtend={use_specextend})...", "yellow"))
    model = EaModel.from_pretrained(
        base_model_path   = BASE_MODEL,
        ea_model_path     = DRAFT_MODEL,
        torch_dtype       = "auto",
        low_cpu_mem_usage = True,
        device_map        = "auto",
    ).eval()

    tokenizer   = model.get_tokenizer()
    accelerator = Accelerator()
    model, tokenizer = accelerator.prepare(model, tokenizer)

    # Iterate lengths
    for length in sorted_lengths:
        records = list(file_to_data[length])

        USE_CHAT_TEMPLATE = True
        ADD_GENERATION_PROMPT = True
        
        # Keep your original govreport filtering rules
        if DATASET_NAME == "govreport":
            if length > 10000:
                records = records[4:]   # skip first 4 if too long
            if length == 2048 and len(records) > 6:
                del records[6]          # special deletion

        # ── Warmup (exactly as original) ───────────────────────────────────────
        print(colored(f"[len={length}] Warming up GPUs...", "yellow"))
        if records:
            # input_ids = tokenizer.encode(records[0]["text"], return_tensors="pt", add_special_tokens=True).to(accelerator.device)
            input_ids = make_input_ids(
                tokenizer, records[0]["text"], accelerator.device,
                use_chat_template=USE_CHAT_TEMPLATE,
                add_generation_prompt=ADD_GENERATION_PROMPT
            )
            
            for _ in range(WARMUP_REPEATS):
                _ = model.eagenerate(
                    input_ids,
                    temperature            = 0.0,
                    max_new_tokens         = WARMUP_TOKENS,
                    output_result_line     = False,
                    verbose                = False,
                    use_specextend         = use_specextend,
                    retrieval_chunk_size   = RETRIEVAL_CHUNK_SIZE,
                    retrieve_top_k         = RETRIEVE_TOP_K,
                    retrieve_every_n_steps = RETRIEVE_EVERY_N_STEPS,
                    retrieval_verbose      = False,
                )
        print(colored("Warmup complete!", "yellow"))

        # ── Accumulators per length (same as original) ────────────────────────
        acc_accepts = []
        total_gen   = 0
        total_time  = 0.0

        # Two replicates per record, checkpoint after each (kept)
        for rec in tqdm(records, desc=f"len={length}"):
            for _ in range(REPLICATES_PER_RECORD):
                # input_ids = tokenizer.encode(rec["text"], return_tensors="pt", add_special_tokens=True).to(accelerator.device)
                input_ids = make_input_ids(
                    tokenizer, rec["text"], accelerator.device,
                    use_chat_template=USE_CHAT_TEMPLATE,
                    add_generation_prompt=ADD_GENERATION_PROMPT
                )
                
                results = model.eagenerate(
                    input_ids,
                    temperature            = 0.0,
                    max_new_tokens         = MAX_NEW_TOKENS_RUN,
                    output_result_line     = True,
                    verbose                = True,
                    use_specextend         = use_specextend,
                    retrieval_chunk_size   = RETRIEVAL_CHUNK_SIZE,
                    retrieve_top_k         = RETRIEVE_TOP_K,
                    retrieve_every_n_steps = RETRIEVE_EVERY_N_STEPS,
                    retrieval_verbose      = False,
                )

                acc_accepts.append(results["avg_accept_length"])
                total_gen  += results["total_generated"]
                total_time += results["inference_time"]

                avg_accept = round(float(np.mean(acc_accepts)), 3)
                tokens_sec = round(total_gen / total_time, 3) if total_time > 1e-9 else 0.0

                summary[length] = {
                    "avg_accept_length":    avg_accept,
                    "tokens_per_sec":       tokens_sec,
                    "total_generated":      total_gen,
                    "total_inference_time": round(total_time, 3),
                    "num_runs":             len(acc_accepts),
                }

                with open(out_path, "w") as f:
                    json.dump({"exp_config": exp_config, "results": summary}, f, indent=2)

    print(colored(f"Experiment complete (SpecExtend={use_specextend})! Results -> {out_path}", "green"))

# ─── RUN BOTH EXPERIMENTS ─────────────────────────────────────────────────────
# run_experiment(use_specextend=False)
run_experiment(use_specextend=True)