Module model

Expand source code
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import json
from typing import Dict
import os
import torch
from torch import nn, Tensor

import logging
logging.basicConfig(
    format=logging.BASIC_FORMAT,
    datefmt='%Y-%m-%d %H:%M:%S',
    level=logging.INFO
)
logger = logging.getLogger(__name__)

class GPTSingleHead(nn.Module):
    """
    Different from directly using GPT2LMHeadModel, this wraps up GPT2LMHeadModel as well as GPT2Tokenizer
    """
    def __init__(self, model_name_or_path: str, max_seq_length: int = 256, do_lower_case: bool = False,
                 special_words_to_add=None):
        super(GPTSingleHead, self).__init__()
        self.config_keys = ['max_seq_length', 'do_lower_case']
        self.do_lower_case = do_lower_case
        if max_seq_length > 1024:
            logging.warning(
                "GPT only allows a max_seq_length of 1024. Value will be set to 1024")
            max_seq_length = 1024
        self.max_seq_length = max_seq_length
        self.gpt = GPT2LMHeadModel.from_pretrained(model_name_or_path)
        self.tokenizer = GPT2Tokenizer.from_pretrained(model_name_or_path, do_lower_case=do_lower_case)
        if special_words_to_add != None:
            self.add_special_words(special_words_to_add)

        self.bos_token_id=self.tokenizer.bos_token_id
        self.eos_token_id=self.tokenizer.eos_token_id
        # self.pad_token_id=self.tokenizer.pad_token_id

    def tokenize(self, text: str):  # default for cls
        return self.tokenizer.convert_tokens_to_ids(self.tokenizer.tokenize(text))

    def add_special_words(self, special_words_to_add):
        orig_num_tokens = len(self.tokenizer)
        num_added_tokens = self.tokenizer.add_special_tokens(special_words_to_add)
        if num_added_tokens > 0:
            self.gpt.resize_token_embeddings(new_num_tokens=orig_num_tokens + num_added_tokens)

    def forward(self, input: Dict[str, torch.Tensor]):
        loss, logits=self.gpt(input["input_ids"],labels=input["input_ids"])[:2]
        return loss, logits

    def get_config_dict(self):
        return {key: self.__dict__[key] for key in self.config_keys}

    def padding_features(self, features_dict_list):
        """
        padding features for a batch
        :param features_dict_list: i.e., batch
        :return: padded batch features
        """
        max_input_len_this_batch = 0

        batch_features = {feature_name: [] for feature_name in features_dict_list[0]}
        for feature_dict in features_dict_list:
            for feature_name, feature_ids in feature_dict.items():
                if feature_name == "input_ids" and len(feature_ids) > max_input_len_this_batch:
                    max_input_len_this_batch = len(feature_ids)
                batch_features[feature_name].append(feature_ids)

        padded_batch_features = {feature_name: [] for feature_name in features_dict_list[0]}
        for feature_name, batch_ids in batch_features.items():

            for each_ids in batch_ids:
                padded = each_ids + [self.tokenizer.pad_token_id] * (max_input_len_this_batch - len(each_ids))
                padded_batch_features[feature_name].append(padded)

        for feature_name, ids in padded_batch_features.items():
            padded_batch_features[feature_name] = torch.tensor(ids)

        return padded_batch_features

    def get_embedding_dimension(self) -> int:
        return self.gpt.config.hidden_size

    def get_config(self) -> int:
        return self.gpt.config

    def save(self, output_path: str):
        self.gpt.save_pretrained(output_path)
        self.tokenizer.save_pretrained(output_path)
        with open(os.path.join(output_path, 'gpt_sh_config.json'), 'w') as f:
            json.dump(self.get_config_dict(), f, indent=2)

    def reload(self, input_path: str):
        """reload from checkpoint weights"""
        return GPTSingleHead.load(input_path + "/0_GPTSingleHead")

    @staticmethod
    def load(input_path: str):
        if not os.path.isfile(os.path.join(input_path, 'gpt_sh_config.json')):
            raise ValueError("In the model path does not find gpt_sh_config.json file, you may have not trained yet")
        with open(os.path.join(input_path, 'gpt_sh_config.json')) as f:
            config = json.load(f)
        return GPTSingleHead(model_name_or_path=input_path, **config)


class EmptyHeads(nn.Module):
    def __init__(self):
        self.config_keys=[]
        super().__init__()

    def forward(self, input: Dict[str, Tensor]):
        return input

    def get_config_dict(self):
        return {key: self.__dict__[key] for key in self.config_keys}

    def save(self, output_path):
        with open(os.path.join(output_path, 'empty_heads_config.json'), 'w') as f:
            json.dump(self.get_config_dict(), f, indent=2)
        torch.save(self.state_dict(), os.path.join(output_path, 'empty_heads.pt'))

    def load_saved(self, input_path):
        self.load_state_dict(torch.load(os.path.join(input_path, '1_EmptyHeads', 'empty_heads.pt')))

    @staticmethod
    def load(input_path,config):
        if not os.path.isfile(os.path.join(input_path, 'empty_heads_config.json')):
            raise ValueError(
                "In the model path does not find empty_heads_config.json file, you may have not trained yet")

        with open(os.path.join(input_path, 'empty_heads_config.json')) as f:
            config = json.load(f)
        model = EmptyHeads()

        if not os.path.isfile(os.path.join(input_path, 'empty_heads.pt')):
            raise ValueError("In the model path does not find state of file, you need to train and get weights first")

        model.load_state_dict(torch.load(os.path.join(input_path, 'empty_heads.pt')))
        return model

Classes

class EmptyHeads

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(Model, self).__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.

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

Expand source code
class EmptyHeads(nn.Module):
    def __init__(self):
        self.config_keys=[]
        super().__init__()

    def forward(self, input: Dict[str, Tensor]):
        return input

    def get_config_dict(self):
        return {key: self.__dict__[key] for key in self.config_keys}

    def save(self, output_path):
        with open(os.path.join(output_path, 'empty_heads_config.json'), 'w') as f:
            json.dump(self.get_config_dict(), f, indent=2)
        torch.save(self.state_dict(), os.path.join(output_path, 'empty_heads.pt'))

    def load_saved(self, input_path):
        self.load_state_dict(torch.load(os.path.join(input_path, '1_EmptyHeads', 'empty_heads.pt')))

    @staticmethod
    def load(input_path,config):
        if not os.path.isfile(os.path.join(input_path, 'empty_heads_config.json')):
            raise ValueError(
                "In the model path does not find empty_heads_config.json file, you may have not trained yet")

        with open(os.path.join(input_path, 'empty_heads_config.json')) as f:
            config = json.load(f)
        model = EmptyHeads()

        if not os.path.isfile(os.path.join(input_path, 'empty_heads.pt')):
            raise ValueError("In the model path does not find state of file, you need to train and get weights first")

        model.load_state_dict(torch.load(os.path.join(input_path, 'empty_heads.pt')))
        return model

Ancestors

  • torch.nn.modules.module.Module

Static methods

def load(input_path, config)
Expand source code
@staticmethod
def load(input_path,config):
    if not os.path.isfile(os.path.join(input_path, 'empty_heads_config.json')):
        raise ValueError(
            "In the model path does not find empty_heads_config.json file, you may have not trained yet")

    with open(os.path.join(input_path, 'empty_heads_config.json')) as f:
        config = json.load(f)
    model = EmptyHeads()

    if not os.path.isfile(os.path.join(input_path, 'empty_heads.pt')):
        raise ValueError("In the model path does not find state of file, you need to train and get weights first")

    model.load_state_dict(torch.load(os.path.join(input_path, 'empty_heads.pt')))
    return model

Methods

def forward(self, input: Dict[str, torch.Tensor])

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Expand source code
def forward(self, input: Dict[str, Tensor]):
    return input
def get_config_dict(self)
Expand source code
def get_config_dict(self):
    return {key: self.__dict__[key] for key in self.config_keys}
def load_saved(self, input_path)
Expand source code
def load_saved(self, input_path):
    self.load_state_dict(torch.load(os.path.join(input_path, '1_EmptyHeads', 'empty_heads.pt')))
def save(self, output_path)
Expand source code
def save(self, output_path):
    with open(os.path.join(output_path, 'empty_heads_config.json'), 'w') as f:
        json.dump(self.get_config_dict(), f, indent=2)
    torch.save(self.state_dict(), os.path.join(output_path, 'empty_heads.pt'))
class GPTSingleHead (model_name_or_path: str, max_seq_length: int = 256, do_lower_case: bool = False, special_words_to_add=None)

Different from directly using GPT2LMHeadModel, this wraps up GPT2LMHeadModel as well as GPT2Tokenizer

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

Expand source code
class GPTSingleHead(nn.Module):
    """
    Different from directly using GPT2LMHeadModel, this wraps up GPT2LMHeadModel as well as GPT2Tokenizer
    """
    def __init__(self, model_name_or_path: str, max_seq_length: int = 256, do_lower_case: bool = False,
                 special_words_to_add=None):
        super(GPTSingleHead, self).__init__()
        self.config_keys = ['max_seq_length', 'do_lower_case']
        self.do_lower_case = do_lower_case
        if max_seq_length > 1024:
            logging.warning(
                "GPT only allows a max_seq_length of 1024. Value will be set to 1024")
            max_seq_length = 1024
        self.max_seq_length = max_seq_length
        self.gpt = GPT2LMHeadModel.from_pretrained(model_name_or_path)
        self.tokenizer = GPT2Tokenizer.from_pretrained(model_name_or_path, do_lower_case=do_lower_case)
        if special_words_to_add != None:
            self.add_special_words(special_words_to_add)

        self.bos_token_id=self.tokenizer.bos_token_id
        self.eos_token_id=self.tokenizer.eos_token_id
        # self.pad_token_id=self.tokenizer.pad_token_id

    def tokenize(self, text: str):  # default for cls
        return self.tokenizer.convert_tokens_to_ids(self.tokenizer.tokenize(text))

    def add_special_words(self, special_words_to_add):
        orig_num_tokens = len(self.tokenizer)
        num_added_tokens = self.tokenizer.add_special_tokens(special_words_to_add)
        if num_added_tokens > 0:
            self.gpt.resize_token_embeddings(new_num_tokens=orig_num_tokens + num_added_tokens)

    def forward(self, input: Dict[str, torch.Tensor]):
        loss, logits=self.gpt(input["input_ids"],labels=input["input_ids"])[:2]
        return loss, logits

    def get_config_dict(self):
        return {key: self.__dict__[key] for key in self.config_keys}

    def padding_features(self, features_dict_list):
        """
        padding features for a batch
        :param features_dict_list: i.e., batch
        :return: padded batch features
        """
        max_input_len_this_batch = 0

        batch_features = {feature_name: [] for feature_name in features_dict_list[0]}
        for feature_dict in features_dict_list:
            for feature_name, feature_ids in feature_dict.items():
                if feature_name == "input_ids" and len(feature_ids) > max_input_len_this_batch:
                    max_input_len_this_batch = len(feature_ids)
                batch_features[feature_name].append(feature_ids)

        padded_batch_features = {feature_name: [] for feature_name in features_dict_list[0]}
        for feature_name, batch_ids in batch_features.items():

            for each_ids in batch_ids:
                padded = each_ids + [self.tokenizer.pad_token_id] * (max_input_len_this_batch - len(each_ids))
                padded_batch_features[feature_name].append(padded)

        for feature_name, ids in padded_batch_features.items():
            padded_batch_features[feature_name] = torch.tensor(ids)

        return padded_batch_features

    def get_embedding_dimension(self) -> int:
        return self.gpt.config.hidden_size

    def get_config(self) -> int:
        return self.gpt.config

    def save(self, output_path: str):
        self.gpt.save_pretrained(output_path)
        self.tokenizer.save_pretrained(output_path)
        with open(os.path.join(output_path, 'gpt_sh_config.json'), 'w') as f:
            json.dump(self.get_config_dict(), f, indent=2)

    def reload(self, input_path: str):
        """reload from checkpoint weights"""
        return GPTSingleHead.load(input_path + "/0_GPTSingleHead")

    @staticmethod
    def load(input_path: str):
        if not os.path.isfile(os.path.join(input_path, 'gpt_sh_config.json')):
            raise ValueError("In the model path does not find gpt_sh_config.json file, you may have not trained yet")
        with open(os.path.join(input_path, 'gpt_sh_config.json')) as f:
            config = json.load(f)
        return GPTSingleHead(model_name_or_path=input_path, **config)

Ancestors

  • torch.nn.modules.module.Module

Static methods

def load(input_path: str)
Expand source code
@staticmethod
def load(input_path: str):
    if not os.path.isfile(os.path.join(input_path, 'gpt_sh_config.json')):
        raise ValueError("In the model path does not find gpt_sh_config.json file, you may have not trained yet")
    with open(os.path.join(input_path, 'gpt_sh_config.json')) as f:
        config = json.load(f)
    return GPTSingleHead(model_name_or_path=input_path, **config)

Methods

def add_special_words(self, special_words_to_add)
Expand source code
def add_special_words(self, special_words_to_add):
    orig_num_tokens = len(self.tokenizer)
    num_added_tokens = self.tokenizer.add_special_tokens(special_words_to_add)
    if num_added_tokens > 0:
        self.gpt.resize_token_embeddings(new_num_tokens=orig_num_tokens + num_added_tokens)
def forward(self, input: Dict[str, torch.Tensor])

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Expand source code
def forward(self, input: Dict[str, torch.Tensor]):
    loss, logits=self.gpt(input["input_ids"],labels=input["input_ids"])[:2]
    return loss, logits
def get_config(self) ‑> int
Expand source code
def get_config(self) -> int:
    return self.gpt.config
def get_config_dict(self)
Expand source code
def get_config_dict(self):
    return {key: self.__dict__[key] for key in self.config_keys}
def get_embedding_dimension(self) ‑> int
Expand source code
def get_embedding_dimension(self) -> int:
    return self.gpt.config.hidden_size
def padding_features(self, features_dict_list)

padding features for a batch :param features_dict_list: i.e., batch :return: padded batch features

Expand source code
def padding_features(self, features_dict_list):
    """
    padding features for a batch
    :param features_dict_list: i.e., batch
    :return: padded batch features
    """
    max_input_len_this_batch = 0

    batch_features = {feature_name: [] for feature_name in features_dict_list[0]}
    for feature_dict in features_dict_list:
        for feature_name, feature_ids in feature_dict.items():
            if feature_name == "input_ids" and len(feature_ids) > max_input_len_this_batch:
                max_input_len_this_batch = len(feature_ids)
            batch_features[feature_name].append(feature_ids)

    padded_batch_features = {feature_name: [] for feature_name in features_dict_list[0]}
    for feature_name, batch_ids in batch_features.items():

        for each_ids in batch_ids:
            padded = each_ids + [self.tokenizer.pad_token_id] * (max_input_len_this_batch - len(each_ids))
            padded_batch_features[feature_name].append(padded)

    for feature_name, ids in padded_batch_features.items():
        padded_batch_features[feature_name] = torch.tensor(ids)

    return padded_batch_features
def reload(self, input_path: str)

reload from checkpoint weights

Expand source code
def reload(self, input_path: str):
    """reload from checkpoint weights"""
    return GPTSingleHead.load(input_path + "/0_GPTSingleHead")
def save(self, output_path: str)
Expand source code
def save(self, output_path: str):
    self.gpt.save_pretrained(output_path)
    self.tokenizer.save_pretrained(output_path)
    with open(os.path.join(output_path, 'gpt_sh_config.json'), 'w') as f:
        json.dump(self.get_config_dict(), f, indent=2)
def tokenize(self, text: str)
Expand source code
def tokenize(self, text: str):  # default for cls
    return self.tokenizer.convert_tokens_to_ids(self.tokenizer.tokenize(text))