"""Exercise the classifier contract on a tiny randomly initialized Gemma4 text model."""
import torch
from transformers import Gemma4TextConfig, Gemma4TextModel
from peft import LoraConfig, TaskType, get_peft_model
from gemma_classifier import GemmaClassifier

torch.set_num_threads(2)
config=Gemma4TextConfig(vocab_size=128,vocab_size_per_layer_input=128,hidden_size=32,
    intermediate_size=64,num_hidden_layers=2,num_attention_heads=2,num_key_value_heads=1,
    head_dim=16,global_head_dim=16,hidden_size_per_layer_input=8,num_kv_shared_layers=0,
    max_position_embeddings=128,sliding_window=32,layer_types=['sliding_attention','full_attention'],
    use_cache=False,pad_token_id=0)
config._attn_implementation='eager'
backbone=Gemma4TextModel(config).to(dtype=torch.bfloat16)
backbone=get_peft_model(backbone,LoraConfig(task_type=TaskType.FEATURE_EXTRACTION,r=2,lora_alpha=4,
    target_modules=['q_proj','v_proj'],lora_dropout=0.))
model=GemmaClassifier(backbone,['care','equality'])
ids=torch.tensor([[2,10,11,0],[2,20,21,22]])
mask=ids.ne(0).long()
result=model(input_ids=ids,attention_mask=mask,labels=torch.tensor([[0.,1.],[1.,0.]]))
assert result.logits.shape==(2,2) and torch.isfinite(result.loss)
result.loss.backward()
assert model.head.weight.grad is not None
assert any(p.grad is not None and p.grad.abs().sum()>0 for n,p in backbone.named_parameters() if 'lora_' in n)
model.eval()
with torch.inference_mode():
    single=model(input_ids=ids[:1,:3],attention_mask=mask[:1,:3]).logits
    padded=model(input_ids=ids[:1],attention_mask=mask[:1]).logits
assert torch.allclose(single,padded,atol=.02), (single,padded)
assert not hasattr(model,'generate')
assert model.backbone.base_model.model.input_modalities==('text',)
print('Tiny Gemma4 contract passed: multi-label logits/loss, head and LoRA gradients, padding equivalence, text-only, no generation')