Module imodelsx.iprompt.prompt_tune

Classes

class PromptTunedModel (args: argparse.Namespace, loss_func: PrefixLoss, model: transformers.modeling_utils.PreTrainedModel, tokenizer: transformers.tokenization_utils.PreTrainedTokenizer, preprefix: str)

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will have their parameters converted too when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Expand source code

class PromptTunedModel(PrefixModel):
    args: argparse.Namespace
    loss_func: PrefixLoss
    model: transformers.PreTrainedModel
    tokenizer: transformers.PreTrainedTokenizer
    prefix_embedding: nn.Parameter
    def __init__(self, args: argparse.Namespace, loss_func: PrefixLoss, model: transformers.PreTrainedModel, tokenizer: transformers.PreTrainedTokenizer, preprefix: str):
        super().__init__(args=args, loss_func=loss_func, model=model, tokenizer=tokenizer, preprefix=preprefix)
        self.prefix_embedding = self.init_continuous_prefix(num_tokens=args.num_learned_tokens)

    def embed_input_ids(self, input_ids: torch.Tensor, prefix_ids: Optional[torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]:
        assert prefix_ids is None, "cannot provide custom prefix IDs for prompt-tuning"
        token_embeddings = self.token_embedding.forward(input_ids)
        return None, torch.cat(
            (self.prefix_embedding.repeat((len(input_ids), 1, 1)), token_embeddings), dim=1
        )

    @property
    def trainable_params(self) -> Iterable[nn.Parameter]:
        return [self.prefix_embedding]
    
    def serialize(self):
        save_dir = self.args.save_dir_unique
        os.makedirs(save_dir, exist_ok=True)
        torch.save(self.prefix_embedding, open(os.path.join(save_dir, 'prefix_embedding.p'), 'wb'))

    def compute_metrics(self) -> Dict[str, Any]:
        return {
            'embs': self.prefix_embedding.detach().cpu().numpy(),
            'grads': self.prefix_embedding.grad.detach().cpu().numpy(),
        }

Ancestors

• PrefixModel
• torch.nn.modules.module.Module
• abc.ABC

Class variables

var args : argparse.Namespace

var loss_func : PrefixLoss

var model : transformers.modeling_utils.PreTrainedModel

var prefix_embedding : torch.nn.parameter.Parameter

var tokenizer : transformers.tokenization_utils.PreTrainedTokenizer

Instance variables

var trainable_params : Iterable[torch.nn.parameter.Parameter]

Expand source code

@property
def trainable_params(self) -> Iterable[nn.Parameter]:
    return [self.prefix_embedding]

Methods

def compute_metrics(self) ‑> Dict[str, Any]

Expand source code

def compute_metrics(self) -> Dict[str, Any]:
    return {
        'embs': self.prefix_embedding.detach().cpu().numpy(),
        'grads': self.prefix_embedding.grad.detach().cpu().numpy(),
    }

Inherited members

• PrefixModel:
  • check_early_stop
  • compute_loss_and_call_backward
  • embed_input_ids
  • forward
  • prepare_batch
  • serialize