caseus fixes

maximegmd · Oct 13, 2023 · 6048973 · 6048973
1 parent f9a831b
commit 6048973
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/src/axolotl/monkeypatch/llama_embeddings_hijack.py b/src/axolotl/monkeypatch/llama_embeddings_hijack.py
@@ -1,3 +1,7 @@
+"""
+patch to add noisy embeddings per https://arxiv.org/abs/2310.05914
+"""
+
 import torch
 import transformers.models.llama.modeling_llama
 from transformers.utils import logging
@@ -6,6 +10,7 @@
 
 
 def replace_llama_embeddings_with_uniform_distribution(noise_alpha=5):
+    # pylint: disable=duplicate-code
     def noised_embed(orig_embed, noise_alpha, model):
         def new_func(input_ids):
             # during training, we add noise to the embedding

diff --git a/src/axolotl/monkeypatch/mistral_embeddings_hijack.py b/src/axolotl/monkeypatch/mistral_embeddings_hijack.py
@@ -1,3 +1,7 @@
+"""
+patch to add noisy embeddings per https://arxiv.org/abs/2310.05914
+"""
+
 import torch
 import transformers.models.mistral.modeling_mistral
 from transformers.utils import logging
@@ -6,6 +10,7 @@
 
 
 def replace_mistral_embeddings_with_uniform_distribution(noise_alpha=5):
+    # pylint: disable=duplicate-code
     def noised_embed(orig_embed, noise_alpha, model):
         def new_func(input_ids):
             # during training, we add noise to the embedding