| from collections import Counter |
| import torch |
| import torch.nn as nn |
|
|
|
|
| class Vocabulary: |
| def __init__(self, min_freq=1): |
| self.min_freq = min_freq |
| self.token2id = {'<PAD>': 0, '<UNK>': 1} |
| self.id2token = {0: '<PAD>', 1: '<UNK>'} |
|
|
| def build(self, texts): |
| counter = Counter(tok for text in texts for tok in text.split()) |
| for token, freq in counter.items(): |
| if freq >= self.min_freq and token not in self.token2id: |
| idx = len(self.token2id) |
| self.token2id[token] = idx |
| self.id2token[idx] = token |
| return self |
|
|
| def encode(self, text, max_len=None): |
| ids = [self.token2id.get(t, 1) for t in text.split()] |
| if max_len: |
| ids = ids[:max_len] |
| return ids if ids else [1] |
|
|
| def __len__(self): |
| return len(self.token2id) |
|
|
|
|
| class AttentionPooling(nn.Module): |
| def __init__(self, hidden_dim): |
| super().__init__() |
| self.attn = nn.Linear(hidden_dim, 1) |
|
|
| def forward(self, h, mask=None): |
| scores = self.attn(h).squeeze(-1) |
| if mask is not None: |
| scores = scores.masked_fill(mask == 0, float('-inf')) |
| weights = torch.softmax(scores, dim=-1) |
| return (h * weights.unsqueeze(-1)).sum(dim=1) |
|
|
|
|
| class BiLSTMClassifier(nn.Module): |
| def __init__(self, vocab_size, embed_dim, hidden_dim, num_layers, num_classes, dropout=0.3): |
| super().__init__() |
| self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0) |
| self.lstm = nn.LSTM( |
| embed_dim, hidden_dim, |
| num_layers=num_layers, |
| batch_first=True, |
| bidirectional=True, |
| dropout=dropout if num_layers > 1 else 0.0, |
| ) |
| self.attention = AttentionPooling(hidden_dim * 2) |
| self.classifier = nn.Sequential( |
| nn.Dropout(dropout), |
| nn.Linear(hidden_dim * 2, num_classes), |
| ) |
|
|
| def forward(self, input_ids, attention_mask=None): |
| x = self.embedding(input_ids) |
| output, _ = self.lstm(x) |
| pooled = self.attention(output, attention_mask) |
| return self.classifier(pooled) |
|
|
