import torch
import torch.nn as nn
from transformers import AutoModel

class TransformerClassifier(nn.Module):
    def __init__(self, model_name, output_dim):
        super(TransformerClassifier, self).__init__()
        self.transformer = AutoModel.from_pretrained(model_name)
        # Freeze bottom 3 layers, unfreeze top layers
        for name, param in self.transformer.named_parameters():
            if "layer.0" in name or "layer.1" in name or "layer.2" in name:
                param.requires_grad = False
        self.fc = nn.Linear(self.transformer.config.hidden_size, output_dim)

    def forward(self, input_ids, attention_mask):
        outputs = self.transformer(input_ids=input_ids, attention_mask=attention_mask)
        hidden_state = outputs.last_hidden_state  # [batch_size, seq_len, hidden_dim]
        pooled_output = hidden_state[:, 0]  # Use CLS token output
        out = self.fc(pooled_output)
        return out