Spaces:
Running on Zero
Running on Zero
File size: 4,617 Bytes
36b4d3a | 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 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | import logging
import os
import time
from dataclasses import dataclass
log = logging.getLogger(__name__)
MAX_WINDOWS = 16 # caps CPU time on very long articles; later text is ignored
WINDOW = 510 # 512 minus the two special tokens
STRIDE = 256
BATCH = 8
@dataclass(frozen=True)
class Spec:
id: str
name: str
repo: str
fake_index: int
trained_on: str
architecture: str
accuracy: float | None = None
titled: bool = False # trained on "<title>..<content>..<end>" inputs
SPECS = [
Spec("argus", "Argus", "hamzab/roberta-fake-news-classification", 0,
"Kaggle Fake and Real News (ISOT)", "RoBERTa-base", 1.0, titled=True),
Spec("sentry", "Sentry", "jy46604790/Fake-News-Bert-Detect", 0,
"40,000+ articles from mixed outlets", "RoBERTa-base"),
Spec("ledger", "Ledger", "winterForestStump/Roberta-fake-news-detector", 0,
"EUvsDisinfo disinformation cases", "RoBERTa-base"),
Spec("prism", "Prism", "dhruvpal/fake-news-bert", 1,
"Mixed real/fake news articles", "DistilBERT", 0.92),
]
def windows(seq: list, size: int, stride: int) -> list[list]:
"""Overlapping windows that always cover the tail of seq."""
last = max(len(seq) - size, 0)
starts = sorted({*range(0, last, stride), last})
return [seq[s : s + size] for s in starts][:MAX_WINDOWS]
def ensemble(p_fakes: list[float]) -> dict | None:
if not p_fakes:
return None
p_fake = sum(p_fakes) / len(p_fakes)
return {
"label": "fake" if p_fake >= 0.5 else "real",
"p_fake": p_fake,
"confidence": abs(p_fake - 0.5) * 2,
}
class Model:
weight = 1.0 # equal votes: the cards report no comparable accuracies
def __init__(self, spec: Spec):
self.spec = spec
self.status = "loading"
self.error = None
def __getattr__(self, name):
return getattr(self.spec, name)
def load(self) -> None:
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.tokenizer = AutoTokenizer.from_pretrained(self.repo)
self.net = AutoModelForSequenceClassification.from_pretrained(self.repo).eval().to(self.device)
def try_load(self) -> None:
try:
self.load()
self.status = "ready"
except Exception as e:
log.warning("model %s unavailable: %s", self.id, e)
self.status, self.error = "unavailable", str(e)
def predict(self, text: str, title: str = "") -> float:
"""Probability the text is fake, averaged over overlapping windows."""
import torch
if self.titled:
text = f"<title>{title}<content>{text}<end>"
ids = self.tokenizer(text, add_special_tokens=False)["input_ids"]
chunks = [self.tokenizer.build_inputs_with_special_tokens(w) for w in windows(ids, WINDOW, STRIDE)]
probs = []
with torch.inference_mode():
for i in range(0, len(chunks), BATCH):
batch = self.tokenizer.pad({"input_ids": chunks[i : i + BATCH]}, return_tensors="pt").to(self.device)
probs += self.net(**batch).logits.softmax(-1)[:, self.fake_index].tolist()
return sum(probs) / len(probs)
def build_models() -> list[Model]:
enabled = os.environ.get("ENABLED_MODELS")
ids = enabled.split(",") if enabled is not None else [s.id for s in SPECS]
return [Model(s) for s in SPECS if s.id in ids]
def load_all(models: list[Model]) -> None:
for m in models:
m.try_load()
def run(models: list[Model], text: str, title: str = "") -> tuple[list[dict], list[str]]:
results, warnings = [], []
for m in models:
if m.status != "ready":
warnings.append(f"{m.name} {'still loading' if m.status == 'loading' else 'unavailable'}")
continue
start = time.perf_counter()
try:
p_fake = m.predict(text, title)
except Exception as e:
log.exception("model %s failed", m.id)
warnings.append(f"{m.name} failed: {e}")
continue
results.append(
{
"id": m.id,
"name": m.name,
"p_fake": p_fake,
"label": "fake" if p_fake >= 0.5 else "real",
"weight": m.weight,
"accuracy": m.accuracy,
"architecture": m.architecture,
"latency_ms": round((time.perf_counter() - start) * 1000),
}
)
return results, warnings
|