Spaces:
Sleeping
Sleeping
File size: 3,979 Bytes
e382dce | 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 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | """BERT-based misinformation classifier (pooler + linear head, Colab/LIAR layout)."""
import os
import numpy as np
import torch
import torch.nn as nn
from transformers import BertModel, BertTokenizer
from src.utils.logger import get_logger
class BERTClassifier(nn.Module):
"""BERT-based misinformation classifier."""
def __init__(self, config=None):
super().__init__()
self.config = config or {}
self.logger = get_logger(__name__)
self.device = "cpu"
self._ready = False
self.bert = None
self.dropout = None
self.classifier = None
self.tokenizer = None
self._try_load()
def _try_load(self):
"""Load BERT model and tokenizer from disk."""
weights_path = os.path.join("models", "bert_classifier.pt")
train_fresh = bool(self.config.get("train_fresh", False))
try:
self.bert = BertModel.from_pretrained("bert-base-uncased")
self.dropout = nn.Dropout(0.3)
self.classifier = nn.Linear(768, 2)
self.tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
if not train_fresh and os.path.exists(weights_path):
state = torch.load(weights_path, map_location=self.device)
self.load_state_dict(state)
self.logger.info("BERT loaded from %s", weights_path)
elif train_fresh:
self.logger.info("BERT train_fresh=True: not loading disk weights")
else:
self.logger.warning("BERT weights not found at %s", weights_path)
self.eval()
self._ready = True
except Exception as e:
self.logger.warning("BERT load failed: %s", e)
self._ready = False
def forward(self, input_ids, attention_mask):
out = self.bert(input_ids=input_ids, attention_mask=attention_mask)
return self.classifier(self.dropout(out.pooler_output))
def load_weights(self, path=None):
"""Load state dict from ``models/bert_classifier.pt`` or given path."""
weights_path = path or os.path.join("models", "bert_classifier.pt")
if not os.path.exists(weights_path):
self.logger.warning("load_weights: missing %s", weights_path)
return
try:
dev = next(self.parameters()).device
except StopIteration:
dev = torch.device("cpu")
state = torch.load(weights_path, map_location=dev)
self.load_state_dict(state, strict=True)
self.eval()
self._ready = True
self.logger.info("BERT weights loaded from %s", weights_path)
def predict(self, texts):
"""
Return P(misinfo) for each text.
Args:
texts: list of strings
Returns:
numpy array of floats between 0 and 1
"""
if not self._ready or self.tokenizer is None:
self.logger.warning("BERT not ready - returning neutral 0.5")
return np.full(len(list(texts)), 0.5)
texts = [str(t) for t in texts]
try:
enc = self.tokenizer(
texts,
padding=True,
truncation=True,
max_length=128,
return_tensors="pt",
)
dev = next(self.parameters()).device
ids = enc["input_ids"].to(dev)
mask = enc["attention_mask"].to(dev)
with torch.no_grad():
logits = self.forward(ids, mask)
probs = torch.softmax(logits, dim=1)
return probs[:, 1].cpu().numpy()
except Exception as e:
self.logger.warning("BERT predict failed: %s", e)
return np.full(len(texts), 0.5)
def predict_proba(self, texts):
"""Return (n, 2) array of [P(credible), P(misinfo)]."""
p_misinfo = self.predict(texts)
return np.column_stack([1 - p_misinfo, p_misinfo])
|