diff --git a/__init__.py b/__init__.py index 0659283..717a5fc 100644 --- a/__init__.py +++ b/__init__.py @@ -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", diff --git a/docs/images/lora_distilbert_classification.png b/docs/images/lora_distilbert_classification.png new file mode 100644 index 0000000..ddef575 Binary files /dev/null and b/docs/images/lora_distilbert_classification.png differ diff --git a/docs/reference/node_reference.md b/docs/reference/node_reference.md index d1369d9..5bba24c 100644 --- a/docs/reference/node_reference.md +++ b/docs/reference/node_reference.md @@ -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. | diff --git a/docs/reference/pt_train_fine_tune_classification_transformer_model.md b/docs/reference/pt_train_fine_tune_classification_transformer_model.md index c0f671c..e30e8d5 100644 --- a/docs/reference/pt_train_fine_tune_classification_transformer_model.md +++ b/docs/reference/pt_train_fine_tune_classification_transformer_model.md @@ -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 diff --git a/docs/reference/ptn_hf_lora_classification_model.md b/docs/reference/ptn_hf_lora_classification_model.md new file mode 100644 index 0000000..7a8e952 --- /dev/null +++ b/docs/reference/ptn_hf_lora_classification_model.md @@ -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 | + +
+Category: PyTorch wrapper - Model + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/examples/workflows/lora_distilbert_classification.json b/examples/workflows/lora_distilbert_classification.json new file mode 100644 index 0000000..b567fb7 --- /dev/null +++ b/examples/workflows/lora_distilbert_classification.json @@ -0,0 +1 @@ +{"last_node_id":36,"last_link_id":52,"nodes":[{"id":19,"type":"PandasShowText","pos":[3504.62841796875,1897.03662109375],"size":[302.7994079589844,112],"flags":{},"order":13,"mode":0,"inputs":[{"name":"text","type":"STRING","widget":{"name":"text"},"link":20}],"outputs":[{"name":"STRING","localized_name":"STRING","type":"STRING","shape":6,"links":null}],"properties":{"Node name for S&R":"PandasShowText"},"widgets_values":["","/container/transfer/models/distilbert_classification.pt"]},{"id":16,"type":"PtnBCEWithLogitsLoss","pos":[2668.492919921875,1800.2408447265625],"size":[315,58],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[{"name":"PTLOSS","localized_name":"PTLOSS","type":"PTLOSS","links":[36],"slot_index":0}],"properties":{"Node name for S&R":"PtnBCEWithLogitsLoss"},"widgets_values":["mean"]},{"id":7,"type":"PtDataLoader","pos":[2293.7060546875,1735.56494140625],"size":[271.68658447265625,173.3039093017578],"flags":{},"order":9,"mode":0,"inputs":[{"name":"dataset","localized_name":"dataset","type":"PTVDATASET","link":24}],"outputs":[{"name":"PTDATALOADER","localized_name":"PTDATALOADER","type":"PTDATALOADER","links":[34],"slot_index":0}],"properties":{"Node name for S&R":"PtDataLoader"},"widgets_values":[32,true,"{\"num_workers\":1}"]},{"id":12,"type":"PtDataLoader","pos":[2290.658447265625,2016.2130126953125],"size":[271.68658447265625,173.3039093017578],"flags":{},"order":10,"mode":0,"inputs":[{"name":"dataset","localized_name":"dataset","type":"PTVDATASET","link":25}],"outputs":[{"name":"PTDATALOADER","localized_name":"PTDATALOADER","type":"PTDATALOADER","links":[35],"slot_index":0}],"properties":{"Node name for S&R":"PtDataLoader"},"widgets_values":[32,false,"{\"num_workers\":1}"]},{"id":21,"type":"PtvHfDatasetWithTokenEncode","pos":[1907.6329345703125,1746.281982421875],"size":[271.5191650390625,178],"flags":{},"order":6,"mode":0,"inputs":[{"name":"encode","localized_name":"encode","type":"PTCALLABLE","link":45}],"outputs":[{"name":"PTVDATASET","localized_name":"PTVDATASET","type":"PTVDATASET","links":[24],"slot_index":0}],"properties":{"Node name for S&R":"PtvHfDatasetWithTokenEncode"},"widgets_values":["imdb","train","text","label",true,true]},{"id":30,"type":"HfTokenizerEncode","pos":[1482.29150390625,1882.735595703125],"size":[315,154],"flags":{},"order":4,"mode":0,"inputs":[{"name":"model_name","type":"STRING","widget":{"name":"model_name"},"link":44}],"outputs":[{"name":"PTCALLABLE","localized_name":"PTCALLABLE","type":"PTCALLABLE","links":[45,46],"slot_index":0}],"properties":{"Node name for S&R":"HfTokenizerEncode"},"widgets_values":["",true,"max_length",true,0]},{"id":24,"type":"PtvHfDatasetWithTokenEncode","pos":[1911.3275146484375,2022.5567626953125],"size":[271.5191650390625,178],"flags":{},"order":7,"mode":0,"inputs":[{"name":"encode","localized_name":"encode","type":"PTCALLABLE","link":46}],"outputs":[{"name":"PTVDATASET","localized_name":"PTVDATASET","type":"PTVDATASET","links":[25],"slot_index":0}],"properties":{"Node name for S&R":"PtvHfDatasetWithTokenEncode"},"widgets_values":["imdb","test","text","label",true,true]},{"id":17,"type":"PtSaveModel","pos":[3495.453125,1735.56494140625],"size":[318.3275146484375,91.27499389648438],"flags":{},"order":12,"mode":0,"inputs":[{"name":"model","localized_name":"model","type":"PTMODEL","link":38}],"outputs":[{"name":"model full path","localized_name":"model full path","type":"STRING","links":[20],"slot_index":0}],"properties":{"Node name for S&R":"PtSaveModel"},"widgets_values":["/container/transfer/models/distilbert_classification.pt"]},{"id":34,"type":"MarkdownNote","pos":[1492.9459228515625,2107.386474609375],"size":[286.51580810546875,101.67364501953125],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["Set max_length to 0 to use model's max_length."],"color":"#432","bgcolor":"#653"},{"id":35,"type":"MarkdownNote","pos":[1910.301025390625,2253.2001953125],"size":[273.0867614746094,101.08975982666016],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["Set encode_return_dict to True for Hugging Face's models."],"color":"#432","bgcolor":"#653"},{"id":31,"type":"CDAStringCreate","pos":[1494.0750732421875,1681.3331298828125],"size":[315,58],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[{"name":"STRING","localized_name":"STRING","type":"STRING","links":[44,50],"slot_index":0}],"properties":{"Node name for S&R":"CDAStringCreate"},"widgets_values":["distilbert-base-uncased"]},{"id":36,"type":"PtnHfLoraClassificationModel","pos":[2100.126708984375,1327.42138671875],"size":[284.3571472167969,186.8283233642578],"flags":{},"order":5,"mode":0,"inputs":[{"name":"model_name","type":"STRING","widget":{"name":"model_name"},"link":50}],"outputs":[{"name":"PTMODEL","localized_name":"PTMODEL","type":"PTMODEL","links":[51,52],"slot_index":0}],"properties":{"Node name for S&R":"PtnHfLoraClassificationModel"},"widgets_values":["",true,0,64,128,0.1]},{"id":29,"type":"PtoAdamW","pos":[2673.749267578125,1107.9647216796875],"size":[285.66900634765625,158.52035522460938],"flags":{},"order":8,"mode":0,"inputs":[{"name":"model","localized_name":"model","type":"PTMODEL","link":51}],"outputs":[{"name":"PTOPTIMIZER","localized_name":"PTOPTIMIZER","type":"PTOPTIMIZER","links":[42],"slot_index":0}],"properties":{"Node name for S&R":"PtoAdamW"},"widgets_values":[0.0003000000000000002,0.9,0.999,0.01,false]},{"id":27,"type":"PtTrainClassificationTransformerModel","pos":[3090.223876953125,1736.2188720703125],"size":[344.3999938964844,292.4398193359375],"flags":{},"order":11,"mode":0,"inputs":[{"name":"model","localized_name":"model","type":"PTMODEL","link":52},{"name":"train_loader","localized_name":"train_loader","type":"PTDATALOADER","link":34},{"name":"optimizer","localized_name":"optimizer","type":"PTOPTIMIZER","link":42},{"name":"loss_function","localized_name":"loss_function","type":"PTLOSS","link":36},{"name":"scheduler","localized_name":"scheduler","type":"PTLRSCHEDULER","shape":7,"link":null},{"name":"val_loader","localized_name":"val_loader","type":"PTDATALOADER","shape":7,"link":35}],"outputs":[{"name":"Model","localized_name":"Model","type":"PTMODEL","links":[38],"slot_index":0},{"name":"train loss","localized_name":"train loss","type":"TENSOR","links":null},{"name":"val loss","localized_name":"val loss","type":"TENSOR","links":null}],"properties":{"Node name for S&R":"PtTrainClassificationTransformerModel"},"widgets_values":[20,true,false,4,false,true]}],"links":[[20,17,0,19,0,"STRING"],[24,21,0,7,0,"PTVDATASET"],[25,24,0,12,0,"PTVDATASET"],[34,7,0,27,1,"PTDATALOADER"],[35,12,0,27,5,"PTDATALOADER"],[36,16,0,27,3,"PTLOSS"],[38,27,0,17,0,"PTMODEL"],[42,29,0,27,2,"PTOPTIMIZER"],[44,31,0,30,0,"STRING"],[45,30,0,21,0,"PTCALLABLE"],[46,30,0,24,0,"PTCALLABLE"],[50,31,0,36,0,"STRING"],[51,36,0,29,0,"PTMODEL"],[52,36,0,27,0,"PTMODEL"]],"groups":[],"config":{},"extra":{"ds":{"scale":1,"offset":[0,0]},"node_versions":{"data-analysis":"68dd090367adffa569c3b3993d56e422f4516d2d","pt-wrapper":"8f87db93c29f15369b8105d9c2e41649768aea05"},"ue_links":[]},"version":0.4} \ No newline at end of file diff --git a/modules/pytorch_wrapper/ptn_hf_lora_classification_model.py b/modules/pytorch_wrapper/ptn_hf_lora_classification_model.py new file mode 100644 index 0000000..6768e1f --- /dev/null +++ b/modules/pytorch_wrapper/ptn_hf_lora_classification_model.py @@ -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,) diff --git a/modules/pytorch_wrapper/ptn_hf_lora_classification_model_def.py b/modules/pytorch_wrapper/ptn_hf_lora_classification_model_def.py new file mode 100644 index 0000000..4882cc1 --- /dev/null +++ b/modules/pytorch_wrapper/ptn_hf_lora_classification_model_def.py @@ -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 diff --git a/pyproject.toml b/pyproject.toml index 323b0eb..0d90451 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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] diff --git a/requirements.txt b/requirements.txt index a7203c7..382e1bd 100644 --- a/requirements.txt +++ b/requirements.txt @@ -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 \ No newline at end of file diff --git a/tests/pytorch_wrapper/ptn_hf_lora_classification_model_def_test.py b/tests/pytorch_wrapper/ptn_hf_lora_classification_model_def_test.py new file mode 100644 index 0000000..721d3b9 --- /dev/null +++ b/tests/pytorch_wrapper/ptn_hf_lora_classification_model_def_test.py @@ -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() diff --git a/tests/pytorch_wrapper/ptn_hf_lora_classification_model_test.py b/tests/pytorch_wrapper/ptn_hf_lora_classification_model_test.py new file mode 100644 index 0000000..ca66fc1 --- /dev/null +++ b/tests/pytorch_wrapper/ptn_hf_lora_classification_model_test.py @@ -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()