Merge branch 'main' of https://github.com/iamgroot42/mimir

iamgroot42 · Feb 11, 2024 · 2773275 · 2773275
2 parents 62664ac + 1fe71b9
commit 2773275
Show file tree

Hide file tree

Showing 8 changed files with 144 additions and 65 deletions.
diff --git a/docs/attacks/blackbox_attacks.html b/docs/attacks/blackbox_attacks.html
@@ -42,7 +42,7 @@ <h1 class="title">Module <code>mimir.attacks.blackbox_attacks</code></h1>
     ZLIB = &#34;zlib&#34; # Done
     MIN_K = &#34;min_k&#34; # Done
     NEIGHBOR = &#34;ne&#34; # Done
-    QUANTILE = &#34;quantile&#34;
+    # QUANTILE = &#34;quantile&#34; # Uncomment when tested implementation is available
 
 
 # Base attack class
@@ -57,7 +57,9 @@ <h1 class="title">Module <code>mimir.attacks.blackbox_attacks</code></h1>
         &#34;&#34;&#34;
         Any attack-specific steps (one-time) preparation
         &#34;&#34;&#34;
-        pass
+        if self.ref_model is not None:
+            self.ref_model.load()
+            self.is_loaded = True
 
     def unload(self):
         if self.ref_model is not None:
@@ -124,7 +126,9 @@ <h2 class="section-title" id="header-classes">Classes</h2>
         &#34;&#34;&#34;
         Any attack-specific steps (one-time) preparation
         &#34;&#34;&#34;
-        pass
+        if self.ref_model is not None:
+            self.ref_model.load()
+            self.is_loaded = True
 
     def unload(self):
         if self.ref_model is not None:
@@ -217,7 +221,9 @@ <h3>Methods</h3>
     &#34;&#34;&#34;
     Any attack-specific steps (one-time) preparation
     &#34;&#34;&#34;
-    pass</code></pre>
+    if self.ref_model is not None:
+        self.ref_model.load()
+        self.is_loaded = True</code></pre>
 </details>
 </dd>
 <dt id="mimir.attacks.blackbox_attacks.Attack.unload"><code class="name flex">
@@ -253,7 +259,7 @@ <h3>Methods</h3>
     ZLIB = &#34;zlib&#34; # Done
     MIN_K = &#34;min_k&#34; # Done
     NEIGHBOR = &#34;ne&#34; # Done
-    QUANTILE = &#34;quantile&#34;</code></pre>
+    # QUANTILE = &#34;quantile&#34; # Uncomment when tested implementation is available</code></pre>
 </details>
 <h3>Ancestors</h3>
 <ul class="hlist">
@@ -274,10 +280,6 @@ <h3>Class variables</h3>
 <dd>
 <div class="desc"></div>
 </dd>
-<dt id="mimir.attacks.blackbox_attacks.BlackBoxAttacks.QUANTILE"><code class="name">var <span class="ident">QUANTILE</span></code></dt>
-<dd>
-<div class="desc"></div>
-</dd>
 <dt id="mimir.attacks.blackbox_attacks.BlackBoxAttacks.REFERENCE_BASED"><code class="name">var <span class="ident">REFERENCE_BASED</span></code></dt>
 <dd>
 <div class="desc"></div>
@@ -314,11 +316,10 @@ <h4><code><a title="mimir.attacks.blackbox_attacks.Attack" href="#mimir.attacks.
 </li>
 <li>
 <h4><code><a title="mimir.attacks.blackbox_attacks.BlackBoxAttacks" href="#mimir.attacks.blackbox_attacks.BlackBoxAttacks">BlackBoxAttacks</a></code></h4>
-<ul class="two-column">
+<ul class="">
 <li><code><a title="mimir.attacks.blackbox_attacks.BlackBoxAttacks.LOSS" href="#mimir.attacks.blackbox_attacks.BlackBoxAttacks.LOSS">LOSS</a></code></li>
 <li><code><a title="mimir.attacks.blackbox_attacks.BlackBoxAttacks.MIN_K" href="#mimir.attacks.blackbox_attacks.BlackBoxAttacks.MIN_K">MIN_K</a></code></li>
 <li><code><a title="mimir.attacks.blackbox_attacks.BlackBoxAttacks.NEIGHBOR" href="#mimir.attacks.blackbox_attacks.BlackBoxAttacks.NEIGHBOR">NEIGHBOR</a></code></li>
-<li><code><a title="mimir.attacks.blackbox_attacks.BlackBoxAttacks.QUANTILE" href="#mimir.attacks.blackbox_attacks.BlackBoxAttacks.QUANTILE">QUANTILE</a></code></li>
 <li><code><a title="mimir.attacks.blackbox_attacks.BlackBoxAttacks.REFERENCE_BASED" href="#mimir.attacks.blackbox_attacks.BlackBoxAttacks.REFERENCE_BASED">REFERENCE_BASED</a></code></li>
 <li><code><a title="mimir.attacks.blackbox_attacks.BlackBoxAttacks.ZLIB" href="#mimir.attacks.blackbox_attacks.BlackBoxAttacks.ZLIB">ZLIB</a></code></li>
 </ul>

diff --git a/docs/attacks/loss.html b/docs/attacks/loss.html
@@ -32,10 +32,13 @@ <h1 class="title">Module <code>mimir.attacks.loss</code></h1>
 &#34;&#34;&#34;
 import torch as ch
 from mimir.attacks.blackbox_attacks import Attack
+from mimir.models import Model
+from mimir.config import ExperimentConfig
 
 
 class LOSSAttack(Attack):
-    def __init__(self, config, model):
+
+    def __init__(self, config: ExperimentConfig, model: Model):
         super().__init__(config, model, ref_model=None)
 
     @ch.no_grad()
@@ -57,7 +60,7 @@ <h2 class="section-title" id="header-classes">Classes</h2>
 <dl>
 <dt id="mimir.attacks.loss.LOSSAttack"><code class="flex name class">
 <span>class <span class="ident">LOSSAttack</span></span>
-<span>(</span><span>config, model)</span>
+<span>(</span><span>config: <a title="mimir.config.ExperimentConfig" href="../config.html#mimir.config.ExperimentConfig">ExperimentConfig</a>, model: <a title="mimir.models.Model" href="../models.html#mimir.models.Model">Model</a>)</span>
 </code></dt>
 <dd>
 <div class="desc"></div>
@@ -66,7 +69,8 @@ <h2 class="section-title" id="header-classes">Classes</h2>
 <span>Expand source code</span>
 </summary>
 <pre><code class="python">class LOSSAttack(Attack):
-    def __init__(self, config, model):
+
+    def __init__(self, config: ExperimentConfig, model: Model):
         super().__init__(config, model, ref_model=None)
 
     @ch.no_grad()

diff --git a/docs/attacks/min_k.html b/docs/attacks/min_k.html
@@ -33,10 +33,13 @@ <h1 class="title">Module <code>mimir.attacks.min_k</code></h1>
 import torch as ch
 import numpy as np
 from mimir.attacks.blackbox_attacks import Attack
+from mimir.models import Model
+from mimir.config import ExperimentConfig
 
 
 class MinKProbAttack(Attack):
-    def __init__(self, config, model):
+
+    def __init__(self, config: ExperimentConfig, model: Model):
         super().__init__(config, model, ref_model=None)
 
     @ch.no_grad()
@@ -75,7 +78,7 @@ <h2 class="section-title" id="header-classes">Classes</h2>
 <dl>
 <dt id="mimir.attacks.min_k.MinKProbAttack"><code class="flex name class">
 <span>class <span class="ident">MinKProbAttack</span></span>
-<span>(</span><span>config, model)</span>
+<span>(</span><span>config: <a title="mimir.config.ExperimentConfig" href="../config.html#mimir.config.ExperimentConfig">ExperimentConfig</a>, model: <a title="mimir.models.Model" href="../models.html#mimir.models.Model">Model</a>)</span>
 </code></dt>
 <dd>
 <div class="desc"></div>
@@ -84,7 +87,8 @@ <h2 class="section-title" id="header-classes">Classes</h2>
 <span>Expand source code</span>
 </summary>
 <pre><code class="python">class MinKProbAttack(Attack):
-    def __init__(self, config, model):
+
+    def __init__(self, config: ExperimentConfig, model: Model):
         super().__init__(config, model, ref_model=None)
 
     @ch.no_grad()

diff --git a/docs/attacks/neighborhood.html b/docs/attacks/neighborhood.html
@@ -42,44 +42,43 @@ <h1 class="title">Module <code>mimir.attacks.neighborhood</code></h1>
 
 from mimir.config import ExperimentConfig
 from mimir.attacks.attack_utils import count_masks, apply_extracted_fills
-from mimir.models import Model
+from mimir.models import Model, ReferenceModel
 from mimir.attacks.blackbox_attacks import Attack
 
 
-# def get_mask_model(config: ExperimentConfig, **kwargs):
-#     if &#34;t5&#34; in config.neighborhood_config.model:
-#         mask_model = T5Model(
-#                 config, model_kwargs=model_kwargs, tokenizer_kwargs=tokenizer_kwargs
-#             )
-#         elif &#34;bert&#34; in config.neighborhood_config.model:
-#             mask_model = BertModel(config)
-#         else:
-#             raise ValueError(f&#34;Unknown model {config.neighborhood_config.model}&#34;)
-
-
 class NeighborhoodAttack(Attack):
+
     def __init__(
         self,
         config: ExperimentConfig,
         target_model: Model,
-        ref_model: Model = None,
+        ref_model: ReferenceModel = None,
         **kwargs,
     ):
         super().__init__(config, target_model, ref_model=None)
         self.ref_model = self._pick_neighbor_model()
         assert issubclass(type(self.ref_model), MaskFillingModel), &#34;ref_model must be MaskFillingModel for neighborhood attack&#34;
 
     def get_mask_model(self):
+        &#34;&#34;&#34;
+            Return the mask filling model.
+        &#34;&#34;&#34;
         return self.ref_model
 
     def create_fill_dictionary(self, data):
+        &#34;&#34;&#34;
+            (Only valid for T5 model) Create fill-fictionary used for random_fills
+        &#34;&#34;&#34;
         neigh_config = self.config.neighborhood_config
         if &#34;t5&#34; in neigh_config.model and neigh_config.random_fills:
             if not self.config.pretokenized:
                 # TODO: maybe can be done if detokenized, but currently not supported
                 self.ref_model.create_fill_dictionary(data)
 
     def _pick_neighbor_model(self):
+        &#34;&#34;&#34;
+            Select and load the mask filling model requested in the config.
+        &#34;&#34;&#34;
         # mask filling t5 model
         mask_model = None
         neigh_config = self.config.neighborhood_config
@@ -129,6 +128,9 @@ <h1 class="title">Module <code>mimir.attacks.neighborhood</code></h1>
         self.ref_model.load()
 
     def get_neighbors(self, documents, **kwargs):
+        &#34;&#34;&#34;
+            Generate neighbors for given documents.
+        &#34;&#34;&#34;
         n_perturbations = kwargs.get(&#34;n_perturbations&#34;, 1)
         span_length = kwargs.get(&#34;span_length&#34;, 10)
         neigh_config = self.config.neighborhood_config
@@ -1145,7 +1147,7 @@ <h3>Inherited members</h3>
 </dd>
 <dt id="mimir.attacks.neighborhood.NeighborhoodAttack"><code class="flex name class">
 <span>class <span class="ident">NeighborhoodAttack</span></span>
-<span>(</span><span>config: <a title="mimir.config.ExperimentConfig" href="../config.html#mimir.config.ExperimentConfig">ExperimentConfig</a>, target_model: <a title="mimir.models.Model" href="../models.html#mimir.models.Model">Model</a>, ref_model: <a title="mimir.models.Model" href="../models.html#mimir.models.Model">Model</a> = None, **kwargs)</span>
+<span>(</span><span>config: <a title="mimir.config.ExperimentConfig" href="../config.html#mimir.config.ExperimentConfig">ExperimentConfig</a>, target_model: <a title="mimir.models.Model" href="../models.html#mimir.models.Model">Model</a>, ref_model: <a title="mimir.models.ReferenceModel" href="../models.html#mimir.models.ReferenceModel">ReferenceModel</a> = None, **kwargs)</span>
 </code></dt>
 <dd>
 <div class="desc"></div>
@@ -1154,28 +1156,38 @@ <h3>Inherited members</h3>
 <span>Expand source code</span>
 </summary>
 <pre><code class="python">class NeighborhoodAttack(Attack):
+
     def __init__(
         self,
         config: ExperimentConfig,
         target_model: Model,
-        ref_model: Model = None,
+        ref_model: ReferenceModel = None,
         **kwargs,
     ):
         super().__init__(config, target_model, ref_model=None)
         self.ref_model = self._pick_neighbor_model()
         assert issubclass(type(self.ref_model), MaskFillingModel), &#34;ref_model must be MaskFillingModel for neighborhood attack&#34;
 
     def get_mask_model(self):
+        &#34;&#34;&#34;
+            Return the mask filling model.
+        &#34;&#34;&#34;
         return self.ref_model
 
     def create_fill_dictionary(self, data):
+        &#34;&#34;&#34;
+            (Only valid for T5 model) Create fill-fictionary used for random_fills
+        &#34;&#34;&#34;
         neigh_config = self.config.neighborhood_config
         if &#34;t5&#34; in neigh_config.model and neigh_config.random_fills:
             if not self.config.pretokenized:
                 # TODO: maybe can be done if detokenized, but currently not supported
                 self.ref_model.create_fill_dictionary(data)
 
     def _pick_neighbor_model(self):
+        &#34;&#34;&#34;
+            Select and load the mask filling model requested in the config.
+        &#34;&#34;&#34;
         # mask filling t5 model
         mask_model = None
         neigh_config = self.config.neighborhood_config
@@ -1225,6 +1237,9 @@ <h3>Inherited members</h3>
         self.ref_model.load()
 
     def get_neighbors(self, documents, **kwargs):
+        &#34;&#34;&#34;
+            Generate neighbors for given documents.
+        &#34;&#34;&#34;
         n_perturbations = kwargs.get(&#34;n_perturbations&#34;, 1)
         span_length = kwargs.get(&#34;span_length&#34;, 10)
         neigh_config = self.config.neighborhood_config
@@ -1271,12 +1286,15 @@ <h3>Methods</h3>
 <span>def <span class="ident">create_fill_dictionary</span></span>(<span>self, data)</span>
 </code></dt>
 <dd>
-<div class="desc"></div>
+<div class="desc"><p>(Only valid for T5 model) Create fill-fictionary used for random_fills</p></div>
 <details class="source">
 <summary>
 <span>Expand source code</span>
 </summary>
 <pre><code class="python">def create_fill_dictionary(self, data):
+    &#34;&#34;&#34;
+        (Only valid for T5 model) Create fill-fictionary used for random_fills
+    &#34;&#34;&#34;
     neigh_config = self.config.neighborhood_config
     if &#34;t5&#34; in neigh_config.model and neigh_config.random_fills:
         if not self.config.pretokenized:
@@ -1288,25 +1306,31 @@ <h3>Methods</h3>
 <span>def <span class="ident">get_mask_model</span></span>(<span>self)</span>
 </code></dt>
 <dd>
-<div class="desc"></div>
+<div class="desc"><p>Return the mask filling model.</p></div>
 <details class="source">
 <summary>
 <span>Expand source code</span>
 </summary>
 <pre><code class="python">def get_mask_model(self):
+    &#34;&#34;&#34;
+        Return the mask filling model.
+    &#34;&#34;&#34;
     return self.ref_model</code></pre>
 </details>
 </dd>
 <dt id="mimir.attacks.neighborhood.NeighborhoodAttack.get_neighbors"><code class="name flex">
 <span>def <span class="ident">get_neighbors</span></span>(<span>self, documents, **kwargs)</span>
 </code></dt>
 <dd>
-<div class="desc"></div>
+<div class="desc"><p>Generate neighbors for given documents.</p></div>
 <details class="source">
 <summary>
 <span>Expand source code</span>
 </summary>
 <pre><code class="python">def get_neighbors(self, documents, **kwargs):
+    &#34;&#34;&#34;
+        Generate neighbors for given documents.
+    &#34;&#34;&#34;
     n_perturbations = kwargs.get(&#34;n_perturbations&#34;, 1)
     span_length = kwargs.get(&#34;span_length&#34;, 10)
     neigh_config = self.config.neighborhood_config

diff --git a/docs/attacks/reference.html b/docs/attacks/reference.html
@@ -31,14 +31,18 @@ <h1 class="title">Module <code>mimir.attacks.reference</code></h1>
     Reference-based attacks.
 &#34;&#34;&#34;
 from mimir.attacks.blackbox_attacks import Attack
+from mimir.models import Model, ReferenceModel
+from mimir.config import ExperimentConfig
 
 
 class ReferenceAttack(Attack):
-    def __init__(self, config, model, reference_model):
-        super().__init__(config, model, reference_model)
 
-    def load(self):
-        self.ref_model.load()
+    def __init__(
+        self, config: ExperimentConfig,
+        model: Model,
+        reference_model: ReferenceModel
+    ):
+        super().__init__(config, model, reference_model)
 
     def _attack(self, document, probs, tokens=None, **kwargs):
         &#34;&#34;&#34;
@@ -62,7 +66,7 @@ <h2 class="section-title" id="header-classes">Classes</h2>
 <dl>
 <dt id="mimir.attacks.reference.ReferenceAttack"><code class="flex name class">
 <span>class <span class="ident">ReferenceAttack</span></span>
-<span>(</span><span>config, model, reference_model)</span>
+<span>(</span><span>config: <a title="mimir.config.ExperimentConfig" href="../config.html#mimir.config.ExperimentConfig">ExperimentConfig</a>, model: <a title="mimir.models.Model" href="../models.html#mimir.models.Model">Model</a>, reference_model: <a title="mimir.models.ReferenceModel" href="../models.html#mimir.models.ReferenceModel">ReferenceModel</a>)</span>
 </code></dt>
 <dd>
 <div class="desc"></div>
@@ -71,11 +75,13 @@ <h2 class="section-title" id="header-classes">Classes</h2>
 <span>Expand source code</span>
 </summary>
 <pre><code class="python">class ReferenceAttack(Attack):
-    def __init__(self, config, model, reference_model):
-        super().__init__(config, model, reference_model)
 
-    def load(self):
-        self.ref_model.load()
+    def __init__(
+        self, config: ExperimentConfig,
+        model: Model,
+        reference_model: ReferenceModel
+    ):
+        super().__init__(config, model, reference_model)
 
     def _attack(self, document, probs, tokens=None, **kwargs):
         &#34;&#34;&#34;