Spaces:
Running on Zero
Running on Zero
| 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 | |
| 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 | |