Spaces:
Sleeping
Sleeping
File size: 3,045 Bytes
469c254 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
import torch
import torch.nn as nn
from transformers import BertModel
from typing import Tuple, Dict
class AttentionLayer(nn.Module):
def __init__(self, hidden_size: int):
super().__init__()
self.attention = nn.Sequential(
nn.Linear(hidden_size, hidden_size),
nn.Tanh(),
nn.Linear(hidden_size, 1)
)
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
attention_weights = torch.softmax(self.attention(x), dim=1)
attended = torch.sum(attention_weights * x, dim=1)
return attended, attention_weights
class HybridFakeNewsDetector(nn.Module):
def __init__(self,
bert_model_name: str = "bert-base-uncased",
lstm_hidden_size: int = 256,
lstm_num_layers: int = 2,
dropout_rate: float = 0.3,
num_classes: int = 2):
super().__init__()
# BERT encoder
self.bert = BertModel.from_pretrained(bert_model_name)
bert_hidden_size = self.bert.config.hidden_size
# BiLSTM layer
self.lstm = nn.LSTM(
input_size=bert_hidden_size,
hidden_size=lstm_hidden_size,
num_layers=lstm_num_layers,
batch_first=True,
bidirectional=True
)
# Attention layer
self.attention = AttentionLayer(lstm_hidden_size * 2)
# Classification head
self.classifier = nn.Sequential(
nn.Dropout(dropout_rate),
nn.Linear(lstm_hidden_size * 2, lstm_hidden_size),
nn.ReLU(),
nn.Dropout(dropout_rate),
nn.Linear(lstm_hidden_size, num_classes)
)
def forward(self, input_ids: torch.Tensor,
attention_mask: torch.Tensor) -> Dict[str, torch.Tensor]:
# Get BERT embeddings
bert_outputs = self.bert(
input_ids=input_ids,
attention_mask=attention_mask
)
bert_embeddings = bert_outputs.last_hidden_state
# Process through BiLSTM
lstm_output, _ = self.lstm(bert_embeddings)
# Apply attention
attended, attention_weights = self.attention(lstm_output)
# Classification
logits = self.classifier(attended)
return {
'logits': logits,
'attention_weights': attention_weights
}
def predict(self, input_ids: torch.Tensor,
attention_mask: torch.Tensor) -> torch.Tensor:
"""Get model predictions."""
outputs = self.forward(input_ids, attention_mask)
return torch.softmax(outputs['logits'], dim=1)
def get_attention_weights(self, input_ids: torch.Tensor,
attention_mask: torch.Tensor) -> torch.Tensor:
"""Get attention weights for interpretability."""
outputs = self.forward(input_ids, attention_mask)
return outputs['attention_weights'] |