app.py

import os, sys
sys.path.insert(0, os.path.dirname(__file__))

import yaml
import torch
import gradio as gr
from huggingface_hub import hf_hub_download, list_repo_files
from transformers import T5Tokenizer, T5EncoderModel

from modules.model import ELF_models
from configs.config import Config, SamplingConfig
from utils.generation_utils import (
    _generate_samples_single_batch,
    _dlm_decode_batch,
    mask_after_eos,
    shift_left,
)
from utils.sampling_utils import get_sampling_steps

# ── Device ───────────────────────────────────────────────
DEVICE    = torch.device("cuda" if torch.cuda.is_available() else "cpu")
HF_REPO   = "embedded-language-flows/ELF-B-xsum-torch"
T5_DIM    = 512   # t5-small encoder dim
VOCAB     = 32100

# ── Config ────────────────────────────────────────────────
config_path = hf_hub_download(HF_REPO, "config.yml")
with open(config_path) as f:
    cfg_dict = yaml.safe_load(f)

config = Config()
for k, v in cfg_dict.items():
    if hasattr(config, k):
        setattr(config, k, v)
config.use_bf16 = torch.cuda.is_available()

# ── Model ─────────────────────────────────────────────────
print("Building ELF-B model...")
model = ELF_models["ELF-B"](
    text_encoder_dim          = T5_DIM,
    max_length                = config.max_length,
    bottleneck_dim            = config.bottleneck_dim,
    num_time_tokens           = config.num_time_tokens,
    num_self_cond_cfg_tokens  = config.num_self_cond_cfg_tokens,
    num_model_mode_tokens     = config.num_model_mode_tokens,
    vocab_size                = VOCAB,
    attn_drop                 = getattr(config, "attn_dropout", 0.0),
    proj_drop                 = getattr(config, "proj_dropout", 0.0),
).to(DEVICE)
model.eval()

# ── Checkpoint — find filename dynamically ────────────────
print("Finding checkpoint...")
ckpt_file = next(
    f for f in list_repo_files(HF_REPO)
    if f.startswith("checkpoint")
)
print(f"Loading {ckpt_file}...")
ckpt_path = hf_hub_download(HF_REPO, ckpt_file)
ckpt      = torch.load(ckpt_path, map_location=DEVICE, weights_only=False)
weights   = ckpt.get("ema_params1", ckpt["params"])
model.load_state_dict(weights)
print("Model ready.")

# ── T5 Encoder (frozen) ───────────────────────────────────
print("Loading T5 encoder...")
tokenizer   = T5Tokenizer.from_pretrained("t5-small", legacy=False)
t5_encoder  = T5EncoderModel.from_pretrained("t5-small").to(DEVICE)
t5_encoder.eval()
for p in t5_encoder.parameters():
    p.requires_grad = False

eos_id = tokenizer.eos_token_id
pad_id = tokenizer.pad_token_id or 0

generator = torch.Generator(device="cpu")

# ── Encode article with T5 ────────────────────────────────
@torch.no_grad()
def encode_article(text: str):
    enc = tokenizer(
        text,
        max_length     = config.max_input_length,
        truncation     = True,
        padding        = "max_length",
        return_tensors = "pt",
    ).to(DEVICE)
    emb = t5_encoder(
        input_ids      = enc["input_ids"],
        attention_mask = enc["attention_mask"],
    ).last_hidden_state                              # (1, 1024, 512)
    emb = (emb - config.latent_mean) / config.latent_std

    # Pad to full max_length (1088) — the model expects cond_seq same size as z
    pad_len = config.max_length - config.max_input_length   # 64
    dtype   = next(model.parameters()).dtype
    emb     = torch.cat([
        emb.to(dtype),
        torch.zeros(1, pad_len, T5_DIM, device=DEVICE, dtype=dtype)
    ], dim=1)                                        # (1, 1088, 512)

    mask = torch.cat([
        enc["attention_mask"].float(),
        torch.zeros(1, pad_len, device=DEVICE)
    ], dim=1)                                        # (1, 1088)

    return emb, mask


# ── Inference ─────────────────────────────────────────────
@torch.no_grad()
def summarise(article: str, n_steps: int, cfg_scale: float, sc_cfg: float):
    if not article.strip():
        return "Please paste an article above."

    cond_seq, cond_mask = encode_article(article)    # (1, 1024, 512), (1, 1024)

    sampling_config = SamplingConfig(
        sampling_method       = "ode",
        num_sampling_steps    = [n_steps],
        cfgs                  = [cfg_scale],
        self_cond_cfg_scales  = [sc_cfg],
        time_schedule         = config.time_schedule,
        sde_gamma             = 0.0,
    )

    dtype = next(model.parameters()).dtype
    t_steps = get_sampling_steps(
        n_steps       = n_steps,
        time_schedule = config.time_schedule,
        P_mean        = config.denoiser_p_mean,
        P_std         = config.denoiser_p_std,
        device        = DEVICE,
        dtype         = dtype,
    )

    d_model = model.text_encoder_dim
    z = (
        torch.randn(1, config.max_length, d_model, generator=generator)
        * config.denoiser_noise_scale
    ).to(DEVICE, dtype=dtype)

    latent = _generate_samples_single_batch(
        model             = model,
        generator         = generator,
        z                 = z,
        t_steps           = t_steps,
        cond_seq          = cond_seq,
        cond_seq_mask     = cond_mask,
        config            = config,
        sampling_config   = sampling_config,
        cfg_scale         = cfg_scale,
        self_cond_cfg_scale = sc_cfg,
    )

    t_final_val   = t_steps[-1].item()
    predicted_ids = _dlm_decode_batch(
        z                   = latent,
        model               = model,
        t_final_val         = t_final_val,
        config              = config,
        self_cond_cfg_scale = sc_cfg,
    )

    # Strip the conditioning prefix, keep only generated summary tokens
    gen_length        = config.max_length - config.max_input_length   # 64 tokens
    cond_len          = cond_mask.to(torch.int32).sum(dim=1)          # (1,)
    predicted_ids     = shift_left(predicted_ids, cond_len, pad_id)[:, :gen_length]
    predicted_ids     = mask_after_eos(predicted_ids, eos_id, pad_id)

    summary = tokenizer.decode(
        predicted_ids[0].cpu().numpy(), skip_special_tokens=True
    )
    return summary.strip() if summary.strip() else "(empty — try more steps or higher CFG)"


# ── Gradio UI ─────────────────────────────────────────────
with gr.Blocks(title="ELF Summarisation") as demo:
    gr.Markdown(
        "# ELF: Embedded Language Flows — Summarisation\n"
        "**ELF-B trained on XSum** · Kaiming He et al., MIT 2026 · "
        "[Paper](https://arxiv.org/abs/2605.10938) · "
        "[GitHub](https://github.com/lillian039/ELF)\n\n"
        "Paste a news article, hit **Summarise**. "
        "Generation is non-autoregressive — all tokens in parallel via flow matching."
    )
    with gr.Row():
        with gr.Column(scale=2):
            article_box = gr.Textbox(
                label       = "Article",
                placeholder = "Paste a news article here...",
                lines       = 12,
            )
            with gr.Row():
                n_steps  = gr.Slider(8, 64, value=64, step=8,
                                     label="Denoising Steps")
                cfg      = gr.Slider(1.0, 5.0, value=2.0, step=0.5,
                                     label="CFG Scale")
                sc_cfg   = gr.Slider(1.0, 5.0, value=1.0, step=0.5,
                                     label="Self-Cond CFG")
            btn = gr.Button("Summarise", variant="primary")
        with gr.Column(scale=1):
            summary_box = gr.Textbox(label="Summary", lines=6)

    btn.click(
        fn      = summarise,
        inputs  = [article_box, n_steps, cfg, sc_cfg],
        outputs = summary_box,
    )

    gr.Examples(
        examples=[
            [
                "The European Space Agency has confirmed that its Mars rover "
                "will launch in 2028, two years later than planned due to technical "
                "issues with the drilling system. The rover, named Rosalind Franklin, "
                "is designed to search for signs of past life by drilling up to two "
                "metres below the Martian surface.",
                64, 2.0, 1.0
            ],
        ],
        inputs=[article_box, n_steps, cfg, sc_cfg],
    )

demo.launch()