Fix docstrings

armbues · Nov 26, 2024 · 8a6691c · 8a6691c
1 parent 005e927
commit 8a6691c
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
diff --git a/sillm/experimental/distillation.py b/sillm/experimental/distillation.py
@@ -11,11 +11,12 @@
 
 logger = logging.getLogger("sillm")
 
-########
-# References:
-# Geoffrey Hinton, Oriol Vinyals and Jeff Dean. Distilling the Knowledge in a Neural Network. https://arxiv.org/abs/1503.02531
-########
 class Distillation(TrainableLoRA):
+    """
+    Trainable distillation LLM.
+    References:
+    https://arxiv.org/abs/1503.02531
+    """
     @staticmethod
     def from_model(target_llm: LLM, draft_llm: LLM, **kwargs):
         """

diff --git a/sillm/experimental/logit_filter.py b/sillm/experimental/logit_filter.py
@@ -4,6 +4,9 @@
 import mlx.core as mx
 
 class LogitFilter:
+    """
+    Logit mask filter.
+    """
     def __init__(self,
                  tokenizer,
                  output_size: int
@@ -17,7 +20,7 @@ def reset(self):
     def __call__(self,
                  logits: mx.array
                  ) -> mx.array:
-        raise NotImplementedError("Class structure.StructureEnforcer is used for inheritance only")
+        raise NotImplementedError("Class LogitFilter is used for inheritance only")
 
 class ASCIIFilter(LogitFilter):
     """