Why?

It's tiny. 60m parameters and it can be easily fitted into... well anything.

It's insanely easy to lobotomize though, so I've devised a multitude of methods to teach this model new behaviors without squashing it like a bug.

WHY THOUGH!?!?!

Guidance. It's being taught, the rule of 3. It's going to learn, exactly how to guide the image generation in a way that no guidance tool has ever accomplished before.

https://huggingface.co/AbstractPhil/T5-Small-Human-Attentive-Try2-Pass3
Not only that, but... it's working. It's adhering internally to large amounts of data; learning exponential amounts of methods to normalize these captions into solidified rule of 3 depictions based on object association, reason, and more.

The input goes in, the output is starting to diverge from it's echo.

Essentially, I've devised a way to train it's math, without destroying it.

The T5 small is a powerful tool.

You can teach it almost anything in a short period of time, and it will be capable of using that anything with almost no training steps.

If you train it in the wrong way though, it'll fall apart like a house of cards on the roof of a bullet train.

Below is a colab drag and drop - required to set the output repo for uploading to a huggingface repo, or removing the lines that upload to repo - blueprint for training the T5-small's internals in a robust and careful methodology.

I'm essentially DESTROYING it's capability currently. So with this, it's reshaping the internals to conform to the goals of... captions.

The t5 small is trained with 3 forms of translation and 1 command for "summarize: ", which are all inherently incapable of their baseline jobs.

However, I've given it a new purpose; caption. This, is currently defining the rule of 3 pathways that will be trainable into diffusion models soon enough.

This version of the T5-small is different. As it's repurposed, it'll become more and more fluent at these pathways, devising more and more complex methodologies based on it's internals - forming even more accurate guidance than before in more complex methodologies and ways.

Given the drop of adherence and the reduction of loss reaching a certain point, with the multi-prompt response from BLEU reaching a certain peak point; this model will be ready to trained into SDXL as a flat guidance spine lora.

What does it mean though?

I'm teaching the T5 to be a solid guidance structure for any diffusion model; as a supplement to the existing encoders. Like a cybernetic limb, or an ear, or in this case; a part of the brain implanted in a way that guides the vectors in a carefully utilizable way.

It means... it'll guide SDXL in a way that has never been guided before; both speeding up the inference, improving fidelity, quality, and more; all because the baseline SDXL variations will be allowed to focus more of their internal mechanisms ON THEIR OWN GOALS.

It's like having an extra set of eyes helping guide the output - not through a blanket mathematics response, but through an interpolated learned response that can simply be snapped on like a lora.

# train_t5_small_human_attentive_full_pipeline.py

import os, random, csv, collections
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import IterableDataset

from transformers import (
    T5ForConditionalGeneration,
    T5TokenizerFast,
    DataCollatorForSeq2Seq,
    TrainingArguments,
    Trainer,
    TrainerCallback,
    get_cosine_schedule_with_warmup,
)
from huggingface_hub import login, hf_hub_download
from sacrebleu import corpus_bleu

# ──────────────────────────────────────────────────────────
# 🔐 Login to HF Hub
# ──────────────────────────────────────────────────────────
from google.colab import userdata

os.environ["HF_TOKEN"] = userdata.get('HF_TOKEN')
login(token=os.getenv("HF_TOKEN"))

# ──────────────────────────────────────────────────────────
# ⚙️ Model & Tokenizer (+ higher dropout)
# ──────────────────────────────────────────────────────────
model = T5ForConditionalGeneration.from_pretrained("AbstractPhil/T5-Small-Human-Attentive-Try2-Pass2")
tokenizer = T5TokenizerFast.from_pretrained("t5-small")

# bump dropout for heavy‐LR regime (Tier 6)
model.config.dropout_rate = 0.3
model.config.attention_dropout = 0.3

# Task prefix config
model.config.task_specific_params = {
    "caption": {
        "prefix":         "caption: ",
        "max_length":     64,
        "num_beams":      4,
        "length_penalty": 1.0,
        "early_stopping": True,
    }
}

# unfreeze everything
for p in model.parameters(): p.requires_grad = True

# ──────────────────────────────────────────────────────────
# ░ Tier 2 ░ WeightScheduler “slow-bleed”
# ──────────────────────────────────────────────────────────
class WeightScheduler:
    def __init__(self, total_steps, high_start=5.0, low_start=0.2, end=1.0):
        """
        total_steps: number of steps over which to decay extremes → 1.0
        high_start: initial multiplier for 'high' tokens
        low_start:  initial multiplier for 'low' tokens
        end:        final multiplier for both
        """
        self.total_steps = total_steps
        self.high_start, self.low_start, self.end = high_start, low_start, end

    def get_weight(self, step, token_class):
        """
        step: current global step
        token_class: 0=neutral, 1=high, 2=low (scalar or tensor)
        returns: same shape as token_class, float weights
        """
        frac = min(step / self.total_steps, 1.0)
        high_w = self.high_start + frac * (self.end - self.high_start)
        low_w  = self.low_start  + frac * (self.end - self.low_start)
        # branch on class
        if isinstance(token_class, torch.Tensor):
            w = torch.ones_like(token_class, dtype=torch.float, device=token_class.device)
            w = torch.where(token_class == 1, high_w, w)
            w = torch.where(token_class == 2, low_w,  w)
            return w
        else:
            return {0:1.0, 1:high_w, 2:low_w}[int(token_class)]

# ──────────────────────────────────────────────────────────
# ░ Tier 4 ░ MaskRescaler
# ──────────────────────────────────────────────────────────
class MaskRescaler:
    def __init__(self, window=2, factor=1.2):
        """
        window: number of tokens on each side of a masked position
        factor: multiplier for neighbor weights
        """
        self.window, self.factor = window, factor

    def rescale(self, weights, mask_positions):
        """
        weights: [B, L] tensor of per-token weights
        mask_positions: list of lists of masked indices per example
        """
        B, L = weights.shape
        for b, positions in enumerate(mask_positions):
            for idx in positions:
                left  = max(0, idx - self.window)
                right = min(L, idx + self.window + 1)
                weights[b, left:right] *= self.factor
        return weights

# ──────────────────────────────────────────────────────────
# ░ Tier 5 ░ MixedCaptionDataset (mask+half-swap)
# ──────────────────────────────────────────────────────────
class MixedCaptionDataset(IterableDataset):
    def __init__(self,
                 tokenizer,
                 repo_id: str,
                 examples_per_file: int = 5_000_000,
                 num_files: int = 10,
                 batch_size: int = 256,
                 low_q: float = 0.4,      # mid-quantile mask
                 high_q: float = 0.6,
                 decay_step: int = 1000,  # decay masked counts
                 swap_prob: float = 0.5   # Tier 5 mixing
    ):
        self.tokenizer   = tokenizer
        self.repo_id     = repo_id
        self.files       = [f"captions/caption_{i+2:03d}.csv" for i in range(num_files)]
        self.batch_size  = batch_size
        self.low_q, self.high_q, self.decay_step = low_q, high_q, decay_step
        self.swap_prob   = swap_prob
        self.examples    = examples_per_file * num_files

        # running global counts of label‐tokens
        self.token_counts = collections.Counter()

    def __len__(self):
        return self.examples

    def __iter__(self):
        # stream each file, batch‐tokenize, update counts & mask
        for rel in self.files:
            path = hf_hub_download(self.repo_id, rel, repo_type="dataset")
            with open(path, encoding="utf-8") as f:
                reader, buf = csv.DictReader(f), []
                for row in reader:
                    text = row.get("text","").strip()
                    if not text: continue
                    buf.append(("caption: "+text, text))
                    if len(buf) >= self.batch_size:
                        for sample in self._batch_tokenize(buf):
                            yield sample
                        buf.clear()
                if buf:
                    for sample in self._batch_tokenize(buf):
                        yield sample

    def _batch_tokenize(self, batch):
        inputs, targets = zip(*batch)
        enc = self.tokenizer(
            list(inputs),
            padding="max_length", truncation=True, max_length=256,
            return_tensors="pt"
        )
        lbl = self.tokenizer(
            list(targets),
            padding="max_length", truncation=True, max_length=64,
            return_tensors="pt"
        )

        # 1) update global counts
        flat = lbl.input_ids.view(-1).tolist()
        valid = [tok for tok in flat if tok != self.tokenizer.pad_token_id]
        self.token_counts.update(valid)

        # 2) select mid-quantile to mask
        items = sorted(self.token_counts.items(), key=lambda x: x[1])
        n = len(items)
        if n:
            toks, cnts = zip(*items)
            low_i  = int(n * self.low_q)
            high_i = int(n * self.high_q)
            mask_set = set(toks[low_i:high_i])
        else:
            mask_set = set()

        # 3) decay their counts
        for t in mask_set:
            self.token_counts[t] = max(0, self.token_counts[t] - self.decay_step)

        # 4) now build samples, either masked or half-swap
        samples = []
        pad = self.tokenizer.pad_token_id
        for i in range(len(batch)):
            label_ids = lbl.input_ids[i]
            if random.random() < self.swap_prob:
                # we'll handle swapping higher up: mark with special flag
                samples.append({
                    "input_ids":      enc.input_ids[i].tolist(),
                    "attention_mask": enc.attention_mask[i].tolist(),
                    "labels":         lbl.input_ids[i].tolist(),
                    "_swap":          True
                })
            else:
                # mask mid-quantile tokens
                mask_bool = torch.tensor([tok in mask_set for tok in label_ids])
                masked    = torch.where(mask_bool, -100, label_ids)
                samples.append({
                    "input_ids":      enc.input_ids[i].tolist(),
                    "attention_mask": enc.attention_mask[i].tolist(),
                    "labels":         masked.tolist(),
                    "_swap":          False
                })
        # if any swaps, pair them sequentially
        out = []
        i = 0
        while i < len(samples):
            s = samples[i]
            if s["_swap"] and i+1 < len(samples):
                t = samples[i+1]
                # half-swap inputs & labels between s and t
                i_len = len(s["input_ids"])
                l_len = len(s["labels"])
                hi = i_len//2; hl = l_len//2
                i1 = torch.tensor(s["input_ids"]); i2 = torch.tensor(t["input_ids"])
                l1 = torch.tensor(s["labels"]);     l2 = torch.tensor(t["labels"])
                # swap halves
                new_i1 = torch.cat([i1[:hi], i2[hi:]], dim=0)
                new_i2 = torch.cat([i2[:hi], i1[hi:]], dim=0)
                new_l1 = torch.cat([l1[:hl], l2[hl:]], dim=0)
                new_l2 = torch.cat([l2[:hl], l1[hl:]], dim=0)
                for rec, ni, nl in [(s,new_i1,new_l1),(t,new_i2,new_l2)]:
                    rec["input_ids"]      = ni.tolist()
                    rec["attention_mask"] = ((ni!=pad).long()).tolist()
                    rec["labels"]         = nl.tolist()
                out.extend([s,t])
                i += 2
            else:
                out.append(s)
                i += 1
        return out

# ──────────────────────────────────────────────────────────
# ░ Tier 2/3/4 ░ WeightedTrainer
# ──────────────────────────────────────────────────────────
class WeightedTrainer(Trainer):
    def __init__(self, *args,
                 weight_scheduler: WeightScheduler,
                 mask_rescaler:    MaskRescaler,
                 low_q_extreme:  float = 0.1,
                 high_q_extreme: float = 0.9,
                 **kwargs):
        super().__init__(*args, **kwargs)
        # inject our scheduler & rescaler
        self.weight_scheduler = weight_scheduler
        self.mask_rescaler    = mask_rescaler
        self.low_q_extreme    = low_q_extreme
        self.high_q_extreme   = high_q_extreme

        # precompute total steps for scheduler
        if self.args.max_steps > 0:
            total = self.args.max_steps
        else:
            # approximate
            ds_len = len(self.train_dataset)
            bs     = self.args.per_device_train_batch_size * self.args.gradient_accumulation_steps
            total  = int(ds_len/bs) * self.args.num_train_epochs
        self.weight_scheduler.total_steps = total

    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        # 1) Remove labels and our custom swap flag
        labels = inputs.pop("labels")
        inputs.pop("_swap", None)

        # 2) Forward pass
        outputs = model(**inputs, labels=labels)
        logits  = outputs.logits       # shape: [B, L, V]
        B, L, V = logits.size()

        # 3) Flatten labels and select valid positions (ignore_index = -100)
        flat_labels = labels.view(-1)
        valid_mask  = flat_labels != -100
        valid_idx   = valid_mask.nonzero(as_tuple=True)[0]
        valid_lbls  = flat_labels[valid_idx]

        # 4) Get current global step
        step = self.state.global_step

        # ── Tier 2: classify tokens into high/low/neutral based on global counts ──
        gc    = self.train_dataset.token_counts
        items = sorted(gc.items(), key=lambda x: x[1])
        n     = len(items)
        low_i  = int(n * self.low_q_extreme)
        high_i = int(n * self.high_q_extreme)
        low_set  = set(tok for tok, _ in items[:low_i])
        high_set = set(tok for tok, _ in items[high_i:])
        # Build a class vector: 0=neutral, 1=high, 2=low
        classes = torch.zeros_like(flat_labels)
        cls_list = []
        for tid in valid_lbls.tolist():
            if tid in high_set:
                cls_list.append(1)
            elif tid in low_set:
                cls_list.append(2)
            else:
                cls_list.append(0)
        classes[valid_idx] = torch.tensor(cls_list, device=classes.device)

        # 5) Tier 2: compute slow-bleed weights
        bleed_w = self.weight_scheduler.get_weight(step, classes)

        # ── Tier 3: inverse-frequency weighting (per-batch) ──
        uniq, cnts = valid_lbls.unique(return_counts=True)
        inv        = cnts.max().float() / cnts.float()
        inv_freq   = torch.ones(V, device=logits.device)
        inv_freq[uniq] = inv
        inv_w_flat = torch.ones_like(flat_labels, dtype=logits.dtype)
        inv_w_flat[valid_idx] = inv_freq[valid_lbls]

        # 6) Combine bleed + inverse-frequency
        w_flat  = bleed_w * inv_w_flat
        weights = w_flat.view(B, L)

        # ── Tier 4: mask-triggered local rescaling ──
        mask_pos = [(labels[b] == -100).nonzero(as_tuple=True)[0].tolist() for b in range(B)]
        weights  = self.mask_rescaler.rescale(weights, mask_pos)

        # 7) Compute weighted negative log-likelihood loss
        logprobs = F.log_softmax(logits, dim=-1).view(-1, V)
        sel_lp   = logprobs[valid_idx, valid_lbls]
        w_sel    = w_flat[valid_idx]
        loss     = -(w_sel * sel_lp).sum() / w_sel.sum()

        return (loss, outputs) if return_outputs else loss


# ──────────────────────────────────────────────────────────
# ░ Tier 7 ░ BLEU Eval Callback
# ──────────────────────────────────────────────────────────
class BLEUCallback(TrainerCallback):
    def __init__(self, model, tokenizer, references, prompts, every_n_steps=500):
        self.model     = model
        self.tokenizer = tokenizer
        self.refs      = references  # list of list of str
        self.prompts   = prompts     # list of str
        self.every_n   = every_n_steps
        self.metric    = corpus_bleu

    def on_step_end(self, args, state, control, **kwargs):
        if state.global_step and state.global_step % self.every_n == 0:
            self.model.eval()
            preds = []
            for p in self.prompts:
                batch = self.tokenizer(
                    "caption: " + p,
                    truncation=True, padding="max_length",
                    max_length=256, return_tensors="pt"
                ).to(self.model.device)
                out = self.model.generate(
                    **batch,
                    max_length=64,
                    num_beams=4,
                    early_stopping=True,
                )
                dec = self.tokenizer.decode(out[0], skip_special_tokens=True)
                preds.append(dec)
            score = self.metric(preds, [self.refs]).score
            print(f"▶️  Step {state.global_step}: mixed‐caption BLEU = {score:.2f}")
            self.model.train()

# ──────────────────────────────────────────────────────────
# 🧠 Assemble & Launch
# ──────────────────────────────────────────────────────────

# held‐out mix references & prompts for BLEU
# held‐out mix references & prompts for BLEU
mix_prompts = [
    "a room of tacos",
    "a forest of neon mushrooms",
    "a beach covered in colorful seashells",
    "a mountain peak above the clouds",
    "a desk cluttered with vintage cameras",
    "a city skyline at dusk",
    "a cat sleeping on a windowsill",
    "a bouquet of wildflowers in a mason jar",
    "a red sports car on a desert road",
    "a steaming cup of coffee on a book",
    "a group of hot air balloons",
    "a snowy village at night",
    "a winding forest path in autumn",
    "a pair of sneakers by the door",
    "a stack of pancakes with syrup",
    "a woodland stream with stones",
    "a baby elephant playing in mud",
    "a painted guitar leaning on a chair",
    "a row of colorful umbrellas on a beach",
    "a vintage typewriter on a desk",
]

mix_references = [
    "brightly lit room filled entirely with tacos",
    "an enchanted forest glowing under neon mushrooms",
    "a sunlit beach strewn with a rainbow of seashells",
    "a solitary mountain peak rising above a sea of clouds",
    "a wooden desk overflowing with vintage film cameras",
    "a city skyline silhouetted against a dusky purple sky",
    "a fluffy cat curled up on a sunlit windowsill",
    "a mason jar brimming with wildflowers picked from a meadow",
    "a sleek red sports car speeding along an empty desert highway",
    "a steaming cup of coffee resting on an open leather-bound book",
    "a colorful fleet of hot air balloons drifting across a clear sky",
    "a quiet snowy village illuminated by warm street lamps",
    "a winding forest trail blanketed in golden autumn leaves",
    "a worn pair of sneakers casually placed by the wooden door",
    "a stack of fluffy pancakes drenched in maple syrup",
    "a gentle woodland stream flowing over moss-covered stones",
    "a baby elephant joyfully splashing mud with its trunk",
    "a brightly painted guitar leaning against a wooden chair",
    "a line of vibrant umbrellas casting shade on the sandy beach",
    "a retro typewriter sitting on a cluttered writing desk",
]


dataset  = MixedCaptionDataset(
    tokenizer,
    repo_id="AbstractPhil/human-templated-captions-1b",
    examples_per_file=5_000_000,
    num_files=2,
    batch_size=256,
    low_q=0.4, high_q=0.6,
    decay_step=1000,
    swap_prob=0.5,
)
collator = DataCollatorForSeq2Seq(tokenizer, model=model, label_pad_token_id=-100)

# build scheduler + rescaler
# we’ll fill total_steps after Trainer init
dummy_total = 1_000_000
weight_scheduler = WeightScheduler(total_steps=dummy_total,
                                   high_start=5.0, low_start=0.2, end=1.0)
mask_rescaler    = MaskRescaler(window=2, factor=1.2)

training_args = TrainingArguments(
    output_dir                 = "./checkpoints-full",
    per_device_train_batch_size=256,
    gradient_accumulation_steps=4,
    dataloader_num_workers     =6,
    learning_rate              =1e-3,
    optim                      ="adafactor",
    weight_decay               =0.01,
    lr_scheduler_type          ="cosine",
    warmup_steps               =500,
    num_train_epochs           =4,
    logging_steps              =50,
    save_steps                 =500,
    save_total_limit           =5,
    max_grad_norm              =1.0,           
    bf16                       =True,
    remove_unused_columns      =False,
    push_to_hub                =True,
    hub_model_id               ="AbstractPhil/T5-Small-Human-Attentive-Try2-Pass3",
    report_to                  ="none",
)

trainer = WeightedTrainer(
    model            =model,
    args             =training_args,
    train_dataset    =dataset,
    data_collator    =collator,
    callbacks        =[
        BLEUCallback(model, tokenizer, mix_references, mix_prompts, every_n_steps=50)
    ],
    weight_scheduler =weight_scheduler,
    mask_rescaler    =mask_rescaler,
    low_q_extreme    =0.1,
    high_q_extreme   =0.9,
)

if __name__ == "__main__":
    trainer.train()
    trainer.push_to_hub()