| """Development-only joint group encoder. Random head: NOT a usable Select scorer. |
| |
| No downloads, truncation, windowing, generated text or production integration. |
| The caller owns request validation and supplies the provisioned encoder/tokenizer. |
| """ |
| import torch |
| from torch import nn |
|
|
|
|
| def pack_groups(tokenizer, query, groups, cap=8192): |
| """Return one sequence and explicit content ranges in caller order.""" |
| if not 1 <= cap <= 8192 or not query or not 1 <= len(groups) <= 512: |
| raise ValueError("invalid_admission") |
| ids_seen = set() |
| cls, sep = tokenizer.cls_token_id, tokenizer.sep_token_id |
| if cls is None or sep is None: |
| raise ValueError("missing_special_tokens") |
| ids = [cls] + tokenizer.encode(query, add_special_tokens=False, truncation=False) + [sep] |
| ranges = [] |
| for group in groups: |
| gid, text = group['id'], group['text'] |
| if not isinstance(gid, str) or not gid or gid in ids_seen or not isinstance(text, str) or not text: |
| raise ValueError("invalid_group") |
| ids_seen.add(gid) |
| content = tokenizer.encode(text, add_special_tokens=False, truncation=False) |
| if not content: |
| raise ValueError("empty_group_range") |
| start = len(ids) |
| ids.extend(content) |
| ranges.append((start, len(ids))) |
| ids.append(sep) |
| if len(ids) > cap: |
| raise ValueError("token_cap") |
| return ids, ranges |
|
|
|
|
| class JointGroupModel(nn.Module): |
| """One encoder invocation, masked content mean and a shared linear head.""" |
|
|
| def __init__(self, encoder): |
| super().__init__() |
| self.encoder = encoder |
| self.head = nn.Linear(encoder.config.hidden_size, 1) |
|
|
| def forward(self, input_ids, attention_mask, ranges): |
| if input_ids.ndim != 2 or input_ids.shape[0] != 1 or attention_mask.shape != input_ids.shape: |
| raise ValueError("single_sequence_required") |
| previous = 0 |
| if not 1 <= len(ranges) <= 512: |
| raise ValueError("invalid_ranges") |
| for start, end in ranges: |
| if not previous <= start < end <= input_ids.shape[1] or not attention_mask[0, start:end].all(): |
| raise ValueError("invalid_ranges") |
| previous = end |
| positions = torch.arange(input_ids.shape[1], device=input_ids.device).unsqueeze(0) |
| states = self.encoder(input_ids=input_ids, attention_mask=attention_mask, |
| position_ids=positions).last_hidden_state[0] |
| pooled = torch.stack([states[start:end].mean(0) for start, end in ranges]) |
| scores = self.head(pooled).squeeze(-1) |
| if not torch.isfinite(scores).all(): |
| raise ValueError("nonfinite_scores") |
| return scores |
|
|
