Added experimental LoRA training for LLM text classification.

2025-04-01 01:36:11 -07:00
parent 59ed74ba76
commit a063925135
12 changed files with 435 additions and 2 deletions
@@ -166,6 +166,7 @@ from .modules.pytorch_wrapper.ptn_embedding_transformer_linear import PtnEmbeddi
 from .modules.pytorch_wrapper.ptn_gru import PtnGRU
 from .modules.pytorch_wrapper.ptn_gru_linear import PtnGRULinear
 from .modules.pytorch_wrapper.ptn_hf_fine_tuned_classification_model import PtnHfFineTunedClassificationModel
+from .modules.pytorch_wrapper.ptn_hf_lora_classification_model import PtnHfLoraClassificationModel
 from .modules.pytorch_wrapper.ptn_huber_loss import PtnHuberLoss
 from .modules.pytorch_wrapper.ptn_instance_norm2d import PtnInstanceNorm2d
 from .modules.pytorch_wrapper.ptn_kl_div_loss import PtnKLDivLoss
@@ -378,6 +379,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
    "PtnGRU": PtnGRU,
    "PtnGRULinear": PtnGRULinear,
    "PtnHfFineTunedClassificationModel": PtnHfFineTunedClassificationModel,
+    "PtnHfLoraClassificationModel": PtnHfLoraClassificationModel,
    "PtnHuberLoss": PtnHuberLoss,
    "PtnInstanceNorm2d": PtnInstanceNorm2d,
    "PtnKLDivLoss": PtnKLDivLoss,
@@ -590,6 +592,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
    "PtnGRU": "Ptn GRU",
    "PtnGRULinear": "Ptn GRU Linear",
    "PtnHfFineTunedClassificationModel": "Ptn Hf Fine Tuned Classification Model",
+    "PtnHfLoraClassificationModel": "Ptn Hf Lora Classification Model",
    "PtnHuberLoss": "Ptn Huber Loss",
    "PtnInstanceNorm2d": "Ptn Instance Norm 2d",
    "PtnKLDivLoss": "Ptn KL Div Loss",
@@ -161,6 +161,7 @@
 | [Ptn GRU](ptn_gru.md) | A gated recurrent unit (GRU) model consisting of one or more of a recurrent layer.   |
 | [Ptn GRU Linear](ptn_gru_linear.md) | A recurrent neural network (GRU) model with a linear head. |
 | [Ptn Hf Fine Tuned Classification Model](ptn_hf_fine_tuned_classification_model.md) | A binary classification model containing a Hugging Face pretrained Transformer model. |
+| [Ptn Hf Lora Classification Model](ptn_hf_lora_classification_model.md) | A binary classification model containing a Hugging Face pretrained Transformer model. |
 | [Ptn Instance Norm 2d](ptn_instance_norm2d.md) | A normalization model to normalize elements over spatial axes within each channel for each sample. |
 | [Ptn LSTM](ptn_lstm.md) | A long short-term memory  (LSTM) model consisting of one or more of a recurrent layer.   |
 | [Ptn LSTM Linear](ptn_lstm_linear.md) | A recurrent neural network (LSTM) model with a linear head. |
@@ -29,7 +29,6 @@ Fine-tunes a classification Transformer model using a given dataset, loss functi
 | early_stopping_rounds | Int |
 | output_best_val_model | Boolean |
 | classification_metrics | Boolean |
-| scheduler | Ptlrscheduler |
 | val_loader | Ptdataloader |

 ## Output
@@ -0,0 +1,34 @@
+# Ptn Hf Lora Classification Model
+A binary classification model containing a Hugging Face pretrained Transformer model.
+
+    Args:  
+        model_name (str): Name or path of the pretrained Hugging Face model.  
+        use_mean_pooling (bool): If True, average token embeddings using the attention mask.  
+                                Otherwise, use pooler output (if available) or the CLS token.  
+        dropout (float): Dropout rate applied before the final classification layer.  
+                        Set to 0.0 to disable dropout. Must be ≥ 0.  
+        lora_r (int): LoRA rank parameter (dimension of the low-rank matrices).  
+        lora_alpha (int): LoRA scaling factor. The adapted weight is added as:  
+                        (lora_alpha / lora_r) * (A @ B), where  
+                        A is a (m × r) matrix and B is a (r × n) matrix.  
+        lora_dropout (float): Dropout applied to the input of the LoRA layers.
+
+## Input
+| Name | Data type |
+|---|---|
+| model_name | String |
+| use_mean_pooling | Boolean |
+| dropout | Float |
+| lora_r | Int |
+| lora_alpha | Int |
+| lora_dropout | Float |
+
+## Output
+| Data type |
+|---|
+| Ptmodel |
+
+<HR>
+Category: PyTorch wrapper - Model
+
+ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
@@ -0,0 +1,89 @@
+from typing import Any, Dict, Tuple
+import torch
+import torch.nn as nn
+
+from .ptn_hf_lora_classification_model_def import HfLoraClassificationModel
+
+class PtnHfLoraClassificationModel:
+    """
+    A binary classification model containing a Hugging Face pretrained Transformer model.
+
+        Args:  
+            model_name (str): Name or path of the pretrained Hugging Face model.  
+            use_mean_pooling (bool): If True, average token embeddings using the attention mask.  
+                                    Otherwise, use pooler output (if available) or the CLS token.  
+            dropout (float): Dropout rate applied before the final classification layer.  
+                            Set to 0.0 to disable dropout. Must be ≥ 0.  
+            lora_r (int): LoRA rank parameter (dimension of the low-rank matrices).  
+            lora_alpha (int): LoRA scaling factor. The adapted weight is added as:  
+                            (lora_alpha / lora_r) * (A @ B), where  
+                            A is a (m × r) matrix and B is a (r × n) matrix.  
+            lora_dropout (float): Dropout applied to the input of the LoRA layers.  
+
+    category: PyTorch wrapper - Model
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the input types for the function.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required input types.
+        """
+        return {
+            "required": {
+                "model_name": ("STRING", {"default": "", "multiline": False}),
+                "use_mean_pooling": ("BOOLEAN", {"default": True}),
+                "dropout": ("FLOAT", {"default": 0, "min": 0, "max": 1, "step": 0.001}),
+                "lora_r": ("INT", {"default": 8, "min": 1, "max": 1e6}),
+                "lora_alpha": ("INT", {"default": 16, "min": 1, "max": 1e6}),
+                "lora_dropout": ("FLOAT", {"default": 0.1, "min": 0, "max": 1, "step": 0.001}),
+            }
+        }
+
+    RETURN_TYPES: Tuple[str] = ("PTMODEL",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Training"
+
+    @classmethod
+    def IS_CHANGED(cls, **kw):
+        return float("NaN")
+
+    def f(self,
+          model_name: str,
+          use_mean_pooling: bool,
+          dropout: float,
+          lora_r: int = 8,
+          lora_alpha: int = 16,
+          lora_dropout: float = 0.1
+        ) -> Tuple[nn.Module]:
+        """
+        Instantiates a binary classification model containing a Hugging Face pretrained Transformer model.
+
+            Args:
+            model_name (str): Name or path of the pretrained Hugging Face model.
+            use_mean_pooling (bool): If True, average token embeddings using the attention mask.
+            Otherwise, use pooler output or CLS token embedding.
+            dropout (float): Dropout rate applied before the final linear layer.
+                             Set to 0.0 to disable dropout. Must be ≥ 0.    
+            lora_r (int): LoRA rank parameter (dimension of the low-rank matrices).
+            lora_alpha (int): LoRA scaling factor. The adapted weight is added as:
+                            (lora_alpha / lora_r) * (A @ B), where
+                            A is a (m × r) matrix and B is a (r × n) matrix.
+            lora_dropout (float): Dropout applied to the input of the LoRA layers.
+
+            Returns:
+                Tuple[nn.Module]: A tuple containing the instantiated PyTorch model.
+        """
+        with torch.inference_mode(False):
+            model = HfLoraClassificationModel(
+                    model_name=model_name,
+                    use_mean_pooling=use_mean_pooling,
+                    dropout=dropout,
+                    lora_r=lora_r,
+                    lora_alpha=lora_alpha,
+                    lora_dropout=lora_dropout
+                )
+
+        return (model,)
@@ -0,0 +1,175 @@
+from typing import List, Dict, Any, Tuple, Type
+import torch
+import torch.nn as nn
+from transformers import AutoConfig, PreTrainedModel
+from peft import get_peft_model, LoraConfig, TaskType
+from peft.tuners.lora import Linear as LoRALinear
+
+class HfLoraClassificationModel(nn.Module):
+    """Generic transformer-based binary classification model with LoRA adaptation.
+
+    This wrapper supports Hugging Face encoder models fine-tuned via LoRA. It supports 
+    three pooling strategies for sequence classification: mean pooling, pooler output, 
+    or CLS token embedding. Only LoRA-adapted attention layers and the classification 
+    head are trainable; all base model parameters remain frozen.
+
+    Args:
+        model_name (str): Name or path of the pretrained Hugging Face model.
+        use_mean_pooling (bool): If True, average token embeddings using the attention mask.
+                                 Otherwise, use pooler output (if available) or the CLS token.
+        dropout (float): Dropout rate applied before the final classification layer.
+                         Set to 0.0 to disable dropout. Must be ≥ 0.
+        lora_r (int): LoRA rank parameter (dimension of the low-rank matrices).
+        lora_alpha (int): LoRA scaling factor. The adapted weight is added as:
+                          (lora_alpha / lora_r) * (A @ B), where
+                          A is a (m × r) matrix and B is a (r × n) matrix.
+        lora_dropout (float): Dropout applied to the input of the LoRA layers.
+    
+    pragma: skip_doc
+    """
+
+    MODEL_CLASS_PREFIXES: List[Tuple[str, str]] = [
+        ("bert-base-uncased", "BertModel"),
+        ("roberta-base", "RobertaModel"),
+        ("distilbert-base-uncased", "DistilBertModel"),
+        ("albert-base-v2", "AlbertModel"),
+        ("google/electra-base-discriminator", "ElectraModel"),
+    ]
+
+
+    def __init__(self,
+                 model_name: str,
+                 use_mean_pooling: bool = False,
+                 dropout: float = 0.1,
+                 lora_r: int = 8,
+                 lora_alpha: int = 16,
+                 lora_dropout: float = 0.1):
+        super().__init__()
+        if dropout < 0.0:
+            raise ValueError(f"Dropout must be non-negative, got {dropout}")
+
+        self.config = AutoConfig.from_pretrained(model_name)
+        model_class = self._resolve_model_class(model_name)
+        base_model = model_class.from_pretrained(model_name, config=self.config)
+
+        # Apply LoRA
+        target_modules = [
+            # BERT, RoBERTa, ALBERT, Electra
+            "query", "key", "value", "dense",
+
+            # DistilBERT
+            "q_lin", "k_lin", "v_lin", "out_lin",
+
+            # GPT-2
+            "c_attn", "c_proj",
+
+            # LLaMA, Ollama, Visual LLaMA
+            "q_proj", "k_proj", "v_proj", "o_proj",
+
+            # CLIP (Text & Vision Transformers)
+            "attn.q_proj", "attn.k_proj", "attn.v_proj", "attn.out_proj",
+            "self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.out_proj",
+            "proj",  # common for CLIP-ViT, ViT heads
+
+            # T5
+            "q", "k", "v", "o",
+            "wi_0", "wi_1", "wo",
+
+            # ViT (Hugging Face ViTModel, etc.)
+            "attention.query", "attention.key", "attention.value", "attention.output.dense",
+            "mlp.dense_in", "mlp.dense_out"
+
+            # FFN modules
+            "ffn.lin1", "ffn.lin2",                      # DistilBERT
+            "intermediate.dense", "output.dense",       # BERT, RoBERTa, Electra
+            "wi_0", "wi_1", "wo",                        # T5
+            "mlp.dense_in", "mlp.dense_out",            # ViT
+            "mlp.c_fc", "mlp.c_proj",                   # GPT-2
+            "mlp.up_proj", "mlp.down_proj", "mlp.gate_proj"  # LLaMA-family
+        ]
+
+        peft_config = LoraConfig(
+            r=lora_r,
+            lora_alpha=lora_alpha,
+            # target_modules=target_modules,
+            target_modules="all-linear",  # Inject to all linear layers
+            lora_dropout=lora_dropout,
+            bias="lora_only",  # Makes the bias in base model layers trainable only if those layers have LoRA applied.
+            task_type=TaskType.FEATURE_EXTRACTION
+        )
+        self.llm_model = get_peft_model(base_model, peft_config)
+
+        self.lora_injected_layers = []
+        for name, module in self.llm_model.named_modules():
+            if isinstance(module, LoRALinear):
+                self.lora_injected_layers.append(name)
+
+        if not self.lora_injected_layers:
+            raise ValueError("LoRA injected into zero modules — check target_modules.")
+
+        # for name in self.lora_injected_layers:
+        #     print(f"[LoRA] Injected into: {name}")
+
+        self.hidden_size = self.config.hidden_size
+        self.use_mean_pooling = use_mean_pooling
+        self.linear = nn.Linear(self.hidden_size, 1)
+
+        self.use_dropout = dropout > 0.0
+        if self.use_dropout:
+            self.dropout = nn.Dropout(dropout)
+
+        nn.init.xavier_uniform_(self.linear.weight)
+        if self.linear.bias is not None:
+            nn.init.zeros_(self.linear.bias)
+
+        # Freeze base model weights except LoRA and linear
+        for param in self.llm_model.base_model.parameters():
+            param.requires_grad = False
+        for param in self.llm_model.parameters():
+            if param.requires_grad:  # only LoRA params
+                param.requires_grad = True
+        for param in self.linear.parameters():
+            param.requires_grad = True
+
+
+    def _resolve_model_class(self, model_name: str) -> Type[PreTrainedModel]:
+        for prefix, class_name in self.MODEL_CLASS_PREFIXES:
+            if model_name.startswith(prefix):
+                module = __import__("transformers", fromlist=[class_name])
+                return getattr(module, class_name)
+        raise ValueError(f"Model '{model_name}' is not supported.")
+
+    def forward(self, inputs:Tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
+        """
+        Runs a forward pass through the transformer model followed by optional dropout and a classification head.
+
+        Args:
+            inputs (Tuple[torch.Tensor, torch.Tensor]): Tuple of below tensors:
+                input_ids (torch.Tensor): Token IDs of shape (batch_size, seq_len).
+                attention_mask (torch.Tensor): Attention mask of shape (batch_size, seq_len).
+
+        Returns:
+            torch.Tensor: Raw classification logits of shape (batch_size, 1).
+        """
+        input_ids, attention_mask = inputs
+        outputs = self.llm_model(input_ids=input_ids, attention_mask=attention_mask)
+
+        if self.use_mean_pooling:
+            last_hidden = outputs.last_hidden_state
+            # Expand the mask across the hidden dimension
+            # e.g. if hidden size is 768, the (seq_len,) mask becomes (seq_len, 768) per sample.
+            input_mask_expanded = attention_mask.unsqueeze(-1).expand(last_hidden.size())
+            sum_embeddings = torch.sum(last_hidden * input_mask_expanded, dim=1)
+            mask_adjusted_count = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
+            x = sum_embeddings / mask_adjusted_count
+        elif hasattr(outputs, "pooler_output") and outputs.pooler_output is not None:
+            x = outputs.pooler_output
+        else:  # Grab the first feature in the sequence
+            x = outputs.last_hidden_state[:, 0, :]
+
+        assert x.dim() == 2
+
+        if self.use_dropout:
+            x = self.dropout(x)
+        x = self.linear(x)
+        return x
@@ -1,7 +1,7 @@
 [project]
 name = "pt-wrapper"
 description = "PyTorch extension for ComfyUI featuring extensive PyTorch wrapper nodes for seamless tensor operations and PyTorch model training."
-version = "3.5.1"
+version = "3.6.0"
 license = {file = "LICENSE"}

 [project.urls]
@@ -7,3 +7,5 @@ datasets>=3.4.1
 transformers>=4.47.1
 gensim==4.3.2
 sentencepiece>=0.2.0
+peft>=0.15.1
+accelerate>=1.5.2
@@ -0,0 +1,52 @@
+import os
+import sys
+import unittest
+import torch
+
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+from pytorch_wrapper.ptn_hf_lora_classification_model_def import HfLoraClassificationModel
+from hugging_face_wrapper.hf_tokenizer_encode import HfTokenizerEncode
+
+
+class TestHfLoraClassificationModel(unittest.TestCase):
+    def setUp(self):
+        """Setup shared test conditions."""
+        self.text = [
+            "I really liked the scenary.",
+            "I did not like the vision of the creator."
+        ]
+        self.model_names = [
+            "distilbert-base-uncased",
+            "bert-base-uncased",
+            "roberta-base",
+            "albert-base-v2",
+            "google/electra-base-discriminator"
+        ]
+
+    def test_all_models(self):
+        """Run output shape test across all supported models."""
+        for model_name in self.model_names:
+            with self.subTest(model_name=model_name):
+                encode = HfTokenizerEncode().f(
+                    model_name=model_name,
+                    padding=True,
+                    padding_method="max_length",
+                    truncation=True,
+                    max_length=0  # Use model's max length
+                )[0]
+
+                model = HfLoraClassificationModel(
+                    model_name,
+                    use_mean_pooling=True,
+                    dropout=0.1
+                )
+
+                token_dict = encode(self.text)
+                out = model((token_dict["input_ids"], token_dict["attention_mask"]))
+                self.assertEqual(out.size(), torch.Size([2, 1]))
+
+
+if __name__ == "__main__":
+    unittest.main()
@@ -0,0 +1,77 @@
+import os
+import sys
+import unittest
+import torch
+
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+from pytorch_wrapper.ptn_hf_lora_classification_model import PtnHfLoraClassificationModel
+from hugging_face_wrapper.hf_tokenizer_encode import HfTokenizerEncode
+
+
+class TestPtfHfLoraClassificationModel(unittest.TestCase):
+    def setUp(self):
+        """Setup shared test conditions."""
+        self.text = [
+            "I really liked the scenary.",
+            "I did not like the vision of the creator."
+        ]
+        self.model_names = [
+            "distilbert-base-uncased",
+            "bert-base-uncased",
+            "roberta-base",
+            "albert-base-v2",
+            "google/electra-base-discriminator"
+        ]
+
+    def test_all_models(self):
+        """Run output shape test across all supported models."""
+        for model_name in self.model_names:
+            with self.subTest(model_name=model_name):
+                encode = HfTokenizerEncode().f(
+                    model_name=model_name,
+                    padding=True,
+                    padding_method="max_length",
+                    truncation=True,
+                    max_length=0  # Use model's max length
+                )[0]
+
+                model = PtnHfLoraClassificationModel().f(
+                    model_name,
+                    use_mean_pooling=True,
+                    dropout=0.1
+                )[0]
+
+                token_dict = encode(self.text)
+                out = model((token_dict["input_ids"], token_dict["attention_mask"]))
+                self.assertEqual(out.size(), torch.Size([2, 1]))
+
+    def test_all_models_with_full_params(self):
+        """Test with full params for all supported models."""
+        for model_name in self.model_names:
+            with self.subTest(model_name=model_name):
+                encode = HfTokenizerEncode().f(
+                    model_name=model_name,
+                    padding=True,
+                    padding_method="max_length",
+                    truncation=True,
+                    max_length=0  # Use model's max length
+                )[0]
+
+                model = PtnHfLoraClassificationModel().f(
+                    model_name,
+                    use_mean_pooling=True,
+                    dropout=0.1,
+                    lora_r=20,
+                    lora_alpha=10,
+                    lora_dropout=0.2
+                )[0]
+
+                token_dict = encode(self.text)
+                out = model((token_dict["input_ids"], token_dict["attention_mask"]))
+                self.assertEqual(out.size(), torch.Size([2, 1]))
+
+
+if __name__ == "__main__":
+    unittest.main()