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()