OpenBackdoor

1
from .poisoner import Poisoner
2
import torch
3
import torch.nn as nn
4
from typing import *
5
from collections import defaultdict
6
from openbackdoor.utils import logger
7
import random
8

9
class BadNetsPoisoner(Poisoner):
10
    r"""
11
        Poisoner for `BadNets <https://arxiv.org/abs/1708.06733>`_
12
    
13
    Args:
14
        triggers (`List[str]`, optional): The triggers to insert in texts. Default to `['cf', 'mn', 'bb', 'tq']`.
15
        num_triggers (`int`, optional): Number of triggers to insert. Default to 1.
16
    """
17
    def __init__(
18
        self, 
19
        triggers: Optional[List[str]] = ["cf", "mn", "bb", "tq"],
20
        num_triggers: Optional[int] = 1,
21
        **kwargs
22
    ):
23
        super().__init__(**kwargs)
24
        
25
        self.triggers = triggers
26
        self.num_triggers = num_triggers
27
        logger.info("Initializing BadNet poisoner, triggers are {}".format(" ".join(self.triggers)))
28
    
29
    def poison(self, data: list):
30
        poisoned = []
31
        for text, label, poison_label in data:
32
            poisoned.append((self.insert(text), self.target_label, 1))
33
        return poisoned
34

35
    def insert(
36
        self, 
37
        text: str, 
38
    ):
39
        r"""
40
            Insert trigger(s) randomly in a sentence.
41
        
42
        Args:
43
            text (`str`): Sentence to insert trigger(s).
44
        """
45
        words = text.split()
46
        for _ in range(self.num_triggers):
47
            insert_word = random.choice(self.triggers)
48
            position = random.randint(0, len(words))
49
            words.insert(position, insert_word)
50
        return " ".join(words)