import torch, torch.nn as nn, torch.nn.functional as F
from transformers import DistilBertModel

class CNNBranch(nn.Module):
    def __init__(self, vocab=30522, edim=128, nf=128, kernels=(2,3,4)):
        super().__init__()
        self.emb=nn.Embedding(vocab,edim,padding_idx=0)
        self.convs=nn.ModuleList([nn.Conv1d(edim,nf,k,padding=k//2) for k in kernels])
        self.drop=nn.Dropout(0.3); self.out_dim=nf*len(kernels)
    def forward(self,ids):
        x=self.emb(ids).permute(0,2,1)
        return self.drop(torch.cat([F.adaptive_max_pool1d(F.gelu(c(x)),1).squeeze(2) for c in self.convs],1))

class HybridClassifier(nn.Module):
    def __init__(self, n_labels=6, vocab=30522, edim=128, nf=128, kernels=(2,3,4), drop=0.3):
        super().__init__()
        self.bert=DistilBertModel.from_pretrained("distilbert-base-uncased")
        self.cnn=CNNBranch(vocab,edim,nf,kernels)
        fused=768+self.cnn.out_dim
        self.head=nn.Sequential(nn.Dropout(drop),nn.Linear(fused,512),nn.GELU(),nn.Dropout(drop*0.7),nn.Linear(512,n_labels))
        self.n_labels=n_labels
    def forward(self, input_ids, attention_mask):
        cls=self.bert(input_ids=input_ids,attention_mask=attention_mask).last_hidden_state[:,0,:]
        cnn=self.cnn(input_ids)
        return self.head(torch.cat([cls,cnn],1))