File size: 1,611 Bytes
d69f154 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | """The 6.53M-parameter head; BGE-M3 is not embedded in these weights."""
from dataclasses import dataclass
from typing import Optional
import torch
from torch import nn
from transformers import PreTrainedModel
from transformers.utils import ModelOutput
from .configuration_scripture_genome import ScriptureGenomeConfig
@dataclass
class ScriptureGenomeOutput(ModelOutput):
genome: Optional[torch.Tensor] = None
class ScriptureGenomeHeadModel(PreTrainedModel):
config_class = ScriptureGenomeConfig
base_model_prefix = "head"
main_input_name = "embeddings"
def __init__(self, config):
super().__init__(config)
self.head = nn.Sequential(
nn.Linear(config.input_dim, config.hidden_dim),
nn.GELU(),
nn.Linear(config.hidden_dim, config.output_dim),
nn.Sigmoid(),
)
self.post_init()
def forward(self, embeddings, return_dict=None):
"""Accept normalized BGE-M3 CLS embeddings [batch,1024], not token IDs."""
if embeddings.ndim != 2 or embeddings.shape[-1] != self.config.input_dim:
raise ValueError("embeddings must have shape [batch, 1024]")
if not embeddings.is_floating_point() or not torch.isfinite(embeddings).all():
raise ValueError("embeddings must be finite floating-point values")
genome = self.head(embeddings)
use_dict = self.config.use_return_dict if return_dict is None else return_dict
return ScriptureGenomeOutput(genome=genome) if use_dict else (genome,)
ScriptureGenomeHeadModel.register_for_auto_class("AutoModel")
|