Complete a PyTorch pipeline that learns to predict highlight snippets in text sections from start and end positions, then train and validate it.

A SQuAD-style head: two logits per token, for start and end.

import torch
import torch.nn as nn
import torch.nn.functional as F

class SpanModel(nn.Module):
    def __init__(self, encoder, hidden):
        super().__init__()
        self.encoder = encoder
        self.head = nn.Linear(hidden, 2)

    def forward(self, ids, attn):
        out = self.encoder(input_ids=ids,
                           attention_mask=attn)
        s, e = self.head(out.last_hidden_state).unbind(-1)
        neg = torch.finfo(s.dtype).min
        pad = attn == 0
        s = s.masked_fill(pad, neg)
        return s, e.masked_fill(pad, neg)

def train_step(model, opt, b):
    model.train()
    s, e = model(b["ids"], b["attn"])
    loss = (F.cross_entropy(s, b["start"])
            + F.cross_entropy(e, b["end"])) / 2
    opt.zero_grad()
    loss.backward()
    opt.step()
    return loss.item()

Interviewers check that pads are masked, labels are token indices mapped from character offsets, and validation runs under no_grad and eval mode. Decode the best pair with end not before start and a length cap, then score exact match and token F1 against the gold character spans.