41 lines
1.5 KiB
Python
41 lines
1.5 KiB
Python
import sys
|
|
from pathlib import Path
|
|
|
|
# Add the root directory (which contains 'cyberslm') to sys.path
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent))
|
|
|
|
from cyberslm.model.model import CyberSLM as Stage1CyberSLM
|
|
from cyberslm.model.config import CyberSLMConfig
|
|
|
|
class CyberSLM(Stage1CyberSLM):
|
|
"""
|
|
Adapter class to bridge SFT's ModelConfig to Stage 1's CyberSLMConfig,
|
|
allowing us to use the original Stage 1 model architecture verbatim.
|
|
"""
|
|
def __init__(self, cfg):
|
|
# Convert SFT ModelConfig to Stage 1 CyberSLMConfig
|
|
stage1_cfg = CyberSLMConfig(
|
|
vocab_size=cfg.vocab_size,
|
|
hidden_dim=cfg.hidden_size,
|
|
num_layers=cfg.num_layers,
|
|
num_heads=cfg.num_heads,
|
|
head_dim=cfg.head_dim,
|
|
ffn_hidden_dim=cfg.ffn_size,
|
|
max_seq_len=cfg.max_seq_len,
|
|
tie_weights=cfg.weight_tying,
|
|
bias=cfg.bias,
|
|
dropout=cfg.dropout,
|
|
norm_eps=cfg.norm_eps,
|
|
)
|
|
super().__init__(stage1_cfg)
|
|
|
|
def forward(self, input_ids, attention_mask=None, **kwargs):
|
|
"""
|
|
Wrapper around Stage 1's forward pass.
|
|
Stage 1 returns (logits, all_attn_weights); the SFT trainer expects
|
|
just the logits Tensor. The key-padding ``attention_mask`` (1=keep,
|
|
0=pad) is forwarded so right-padded batches do not corrupt real tokens.
|
|
"""
|
|
logits, _ = super().forward(input_ids, attention_mask=attention_mask)
|
|
return logits
|