| import torch
|
| import torch.nn as nn
|
| from transformers import PreTrainedModel, DistilBertConfig, DistilBertModel
|
| from .configuration_dualmedbert import DualMedBertConfig
|
|
|
| class LoRALinear(nn.Module):
|
| def __init__(self, linear: nn.Linear, rank: int = 8, alpha: int = 32):
|
| super().__init__()
|
| self.linear = linear
|
| self.linear.weight.requires_grad_(False)
|
| if self.linear.bias is not None:
|
| self.linear.bias.requires_grad_(False)
|
|
|
| d_in = linear.in_features
|
| d_out = linear.out_features
|
| self.lora_A = nn.Parameter(torch.randn(d_in, rank) * 0.01)
|
| self.lora_B = nn.Parameter(torch.zeros(rank, d_out))
|
| self.scale = alpha / rank
|
|
|
| def forward(self, x):
|
| return self.linear(x) + (x @ self.lora_A @ self.lora_B) * self.scale
|
|
|
| def inject_lora(model, rank: int = 8, alpha: int = 32):
|
| for p in model.parameters():
|
| p.requires_grad_(False)
|
| for layer in getattr(model, "transformer").layer:
|
| layer.attention.q_lin = LoRALinear(layer.attention.q_lin, rank, alpha)
|
| layer.attention.v_lin = LoRALinear(layer.attention.v_lin, rank, alpha)
|
| return model
|
|
|
| class SelfAttnPool(nn.Module):
|
| def __init__(self, hidden: int = 768, heads: int = 4):
|
| super().__init__()
|
| self.heads = heads
|
| self.query = nn.Linear(hidden, heads, bias=False)
|
| self.proj = nn.Linear(hidden * heads, hidden)
|
|
|
| def forward(self, h: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
|
| scores = self.query(h)
|
| mask_3d = mask.unsqueeze(-1).expand_as(scores)
|
| scores = scores.masked_fill(mask_3d == 0, torch.finfo(scores.dtype).min)
|
| weights = torch.softmax(scores, dim=1)
|
| contexts = torch.matmul(weights.transpose(1, 2), h)
|
| flat = contexts.reshape(h.size(0), -1)
|
| return self.proj(flat)
|
|
|
| class DualMedBertForSequenceClassification(PreTrainedModel):
|
| config_class = DualMedBertConfig
|
|
|
| _tied_weights_keys = []
|
|
|
| @property
|
| def all_tied_weights_keys(self):
|
| return {}
|
|
|
|
|
| def __init__(self, config):
|
| super().__init__(config)
|
| self.num_classes = config.num_classes
|
|
|
| base_config = DistilBertConfig()
|
| base = DistilBertModel(base_config)
|
| self.backbone = inject_lora(base, rank=config.lora_rank, alpha=config.lora_alpha)
|
|
|
| self.attn_pool = SelfAttnPool(768, heads=4)
|
| self.fuse = nn.Linear(768 * 2, 768)
|
| self.relu = nn.ReLU()
|
| self.fuse_drop = nn.Dropout(0.3)
|
| self.dropout = nn.Dropout(0.2)
|
| self.head = nn.Linear(768, self.num_classes)
|
|
|
| def _backbone_pass(self, input_ids, attention_mask):
|
| return self.backbone(input_ids=input_ids,
|
| attention_mask=attention_mask).last_hidden_state
|
|
|
| def forward(self, input_ids, attention_mask=None, **kwargs):
|
| if attention_mask is None:
|
| attention_mask = torch.ones_like(input_ids)
|
|
|
| h = self._backbone_pass(input_ids, attention_mask)
|
| cls = h[:, 0, :]
|
| pool = self.attn_pool(h, attention_mask)
|
| x = self.fuse(torch.cat([cls, pool], dim=-1))
|
| x = self.relu(x)
|
| x = self.fuse_drop(x)
|
| x = self.dropout(x)
|
| logits = self.head(x)
|
|
|
| from transformers.modeling_outputs import SequenceClassifierOutput
|
| return SequenceClassifierOutput(
|
| loss=None,
|
| logits=logits,
|
| hidden_states=(h,),
|
| attentions=None
|
| )
|
|
|
