Coverage for transformer_lens/model_bridge/supported_architectures/seed_oss.py: 100%

1"""Seed-OSS architecture adapter. 

2 

3ByteDance's Seed-OSS (``SeedOssForCausalLM``) is a Llama-layout decoder — 

4RMSNorm + RoPE + GQA + gated MLP under identical module paths — with 

5config-gated attention/MLP biases (handled by the shared weight machinery) 

6and no BOS prepending. 

7""" 

8 

9from typing import Any 

10 

11from transformer_lens.model_bridge.supported_architectures.llama import ( 

12 LlamaArchitectureAdapter, 

13) 

14 

15 

16class SeedOssArchitectureAdapter(LlamaArchitectureAdapter): 

17 """Architecture adapter for SeedOssForCausalLM models.""" 

18 

19 def __init__(self, cfg: Any) -> None: 

20 super().__init__(cfg) 

21 # Verified against ByteDance-Seed/Seed-OSS-36B-Instruct's tokenizer. 

22 self.cfg.default_prepend_bos = False 

23 # Seed-OSS ships attention_bias=True with GQA; the Llama parent omits 

24 # bias reshapes, so K/V biases would keep the flat (n_kv*d_head,) layout. 

25 self.weight_processing_conversions = { 

26 **self._qkvo_weight_conversions(include_biases=True), 

27 }