Skip to content

unitorch.models.pegasus¤

PegasusProcessor¤

Bases: HfTextGenerationProcessor

Initializes a PegasusProcessor.

Parameters:

Name Type Description Default
vocab_path str

The path to the vocabulary file.

 required
special_input_ids Dict

Dictionary of special input tokens and their corresponding IDs. Defaults to an empty dictionary.

 dict()
max_seq_length int

The maximum length of input sequences. Defaults to 128.

 128
max_gen_seq_length int

The maximum length of generated sequences. Defaults to 48.

 48
Source code in src/unitorch/models/pegasus/processing.py 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def __init__(
    self,
    vocab_path: str,
    special_input_ids: Optional[Dict] = dict(),
    max_seq_length: Optional[int] = 128,
    max_gen_seq_length: Optional[int] = 48,
):
    """
    Initializes a PegasusProcessor.

    Args:
        vocab_path (str): The path to the vocabulary file.
        special_input_ids (Dict, optional): Dictionary of special input tokens and their corresponding IDs. Defaults to an empty dictionary.
        max_seq_length (int, optional): The maximum length of input sequences. Defaults to 128.
        max_gen_seq_length (int, optional): The maximum length of generated sequences. Defaults to 48.
    """
    tokenizer = get_pegasus_tokenizer(
        vocab_path,
        special_input_ids=special_input_ids,
    )
    tokenizer.bos_token_id = 0
    tokenizer.bos_token = tokenizer.convert_ids_to_tokens(0)
    tokenizer.sep_token = tokenizer.eos_token
    tokenizer.sep_token_id = tokenizer.eos_token_id
    super().__init__(
        tokenizer=tokenizer,
        max_seq_length=max_seq_length,
        max_gen_seq_length=max_gen_seq_length,
    )

PegasusForGeneration¤

Bases: GenericModel

Initializes a PegasusForGeneration model.

Parameters:

Name Type Description Default
config_path str

The path to the model configuration file.

 required
gradient_checkpointing bool

Whether to use gradient checkpointing. Defaults to False.

 False
Source code in src/unitorch/models/pegasus/modeling.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
def __init__(
    self,
    config_path: str,
    gradient_checkpointing: Optional[bool] = False,
):
    """
    Initializes a PegasusForGeneration model.

    Args:
        config_path (str): The path to the model configuration file.
        gradient_checkpointing (bool, optional): Whether to use gradient checkpointing. Defaults to False.
    """
    super().__init__()
    self.config = PegasusConfig.from_json_file(config_path)
    self.config.gradient_checkpointing = gradient_checkpointing
    self.model = PegasusForConditionalGeneration(self.config)

    self.init_weights()

forward ¤

forward(
    input_ids: Tensor,
    attention_mask: Tensor,
    decoder_input_ids: Tensor,
    decoder_attention_mask: Tensor,
)

Performs a forward pass of the PegasusForGeneration model.

Parameters:

Name Type Description Default
input_ids Tensor

Tensor of input token IDs.

 required
attention_mask Tensor

Tensor of attention mask.

 required
decoder_input_ids Tensor

Tensor of decoder input token IDs.

 required
decoder_attention_mask Tensor

Tensor of decoder attention mask.

 required

Returns:

Type Description
Tensor

The model's logits.
Source code in src/unitorch/models/pegasus/modeling.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
def forward(
    self,
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    decoder_input_ids: torch.Tensor,
    decoder_attention_mask: torch.Tensor,
):
    """
    Performs a forward pass of the PegasusForGeneration model.

    Args:
        input_ids (torch.Tensor): Tensor of input token IDs.
        attention_mask (torch.Tensor): Tensor of attention mask.
        decoder_input_ids (torch.Tensor): Tensor of decoder input token IDs.
        decoder_attention_mask (torch.Tensor): Tensor of decoder attention mask.

    Returns:
        (torch.Tensor):The model's logits.
    """
    outputs = self.model(
        input_ids=input_ids,
        attention_mask=attention_mask,
        decoder_input_ids=decoder_input_ids,
        decoder_attention_mask=decoder_attention_mask,
        return_dict=True,
    )
    logits = outputs.logits
    return logits

generate ¤

generate(
    input_ids: Tensor,
    num_beams: Optional[int] = 5,
    decoder_start_token_id: Optional[int] = 0,
    decoder_end_token_id: Optional[
        Union[int, List[int]]
    ] = 1,
    num_return_sequences: Optional[int] = 1,
    min_gen_seq_length: Optional[int] = 0,
    max_gen_seq_length: Optional[int] = 48,
    repetition_penalty: Optional[float] = 1.0,
    no_repeat_ngram_size: Optional[int] = 0,
    early_stopping: Optional[bool] = True,
    length_penalty: Optional[float] = 1.0,
    num_beam_groups: Optional[int] = 1,
    diversity_penalty: Optional[float] = 0.0,
    do_sample: Optional[bool] = False,
    temperature: Optional[float] = 1.0,
    top_k: Optional[int] = 50,
    top_p: Optional[float] = 1.0,
)

Generates sequences using the PegasusForGeneration model.

Parameters:

Name Type Description Default
input_ids Tensor

The input token IDs.

 required
num_beams int

The number of beams for beam search. Defaults to 5.

 5
decoder_start_token_id int

The decoder's start token ID. Defaults to 2.

 0
decoder_end_token_id int or List[int]

The decoder's end token ID. Defaults to 2.

 1
num_return_sequences int

The number of generated sequences to return. Defaults to 1.

 1
min_gen_seq_length int

The minimum length of the generated sequences. Defaults to 0.

 0
max_gen_seq_length int

The maximum length of the generated sequences. Defaults to 48.

 48
repetition_penalty float

The repetition penalty. Defaults to 1.0.

 1.0
no_repeat_ngram_size int

The size of n-grams to avoid repeating. Defaults to 0.

 0
early_stopping bool

Whether to stop generation early. Defaults to True.

 True
length_penalty float

The length penalty. Defaults to 1.0.

 1.0
num_beam_groups int

The number of beam groups for diverse beam search. Defaults to 1.

 1
diversity_penalty float

The diversity penalty. Defaults to 0.0.

 0.0
do_sample bool

Whether to use sampling for generation. Defaults to False.

 False
temperature float

The temperature for sampling. Defaults to 1.0.

 1.0
top_k int

The value for top-k sampling. Defaults to 50.

 50
top_p float

The value for top-p (nucleus) sampling. Defaults to 1.0.

 1.0

Returns:

Name Type Description
GenericOutputs

The generated sequences and their scores.
Source code in src/unitorch/models/pegasus/modeling.py 
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
@torch.no_grad()
def generate(
    self,
    input_ids: torch.Tensor,
    num_beams: Optional[int] = 5,
    decoder_start_token_id: Optional[int] = 0,
    decoder_end_token_id: Optional[Union[int, List[int]]] = 1,
    num_return_sequences: Optional[int] = 1,
    min_gen_seq_length: Optional[int] = 0,
    max_gen_seq_length: Optional[int] = 48,
    repetition_penalty: Optional[float] = 1.0,
    no_repeat_ngram_size: Optional[int] = 0,
    early_stopping: Optional[bool] = True,
    length_penalty: Optional[float] = 1.0,
    num_beam_groups: Optional[int] = 1,
    diversity_penalty: Optional[float] = 0.0,
    do_sample: Optional[bool] = False,
    temperature: Optional[float] = 1.0,
    top_k: Optional[int] = 50,
    top_p: Optional[float] = 1.0,
):
    """
    Generates sequences using the PegasusForGeneration model.

    Args:
        input_ids: The input token IDs.
        num_beams (int, optional): The number of beams for beam search. Defaults to 5.
        decoder_start_token_id (int, optional): The decoder's start token ID. Defaults to 2.
        decoder_end_token_id (int or List[int], optional): The decoder's end token ID. Defaults to 2.
        num_return_sequences (int, optional): The number of generated sequences to return. Defaults to 1.
        min_gen_seq_length (int, optional): The minimum length of the generated sequences. Defaults to 0.
        max_gen_seq_length (int, optional): The maximum length of the generated sequences. Defaults to 48.
        repetition_penalty (float, optional): The repetition penalty. Defaults to 1.0.
        no_repeat_ngram_size (int, optional): The size of n-grams to avoid repeating. Defaults to 0.
        early_stopping (bool, optional): Whether to stop generation early. Defaults to True.
        length_penalty (float, optional): The length penalty. Defaults to 1.0.
        num_beam_groups (int, optional): The number of beam groups for diverse beam search. Defaults to 1.
        diversity_penalty (float, optional): The diversity penalty. Defaults to 0.0.
        do_sample (bool, optional): Whether to use sampling for generation. Defaults to False.
        temperature (float, optional): The temperature for sampling. Defaults to 1.0.
        top_k (int, optional): The value for top-k sampling. Defaults to 50.
        top_p (float, optional): The value for top-p (nucleus) sampling. Defaults to 1.0.

    Returns:
        GenericOutputs: The generated sequences and their scores.
    """
    outputs = self.model.generate(
        input_ids,
        max_length=max_gen_seq_length,
        min_length=min_gen_seq_length,
        num_beams=num_beams,
        do_sample=do_sample,
        decoder_start_token_id=decoder_start_token_id,
        no_repeat_ngram_size=no_repeat_ngram_size,
        early_stopping=early_stopping,
        length_penalty=length_penalty,
        repetition_penalty=repetition_penalty,
        num_return_sequences=num_return_sequences,
        bos_token_id=decoder_start_token_id,
        eos_token_id=decoder_end_token_id,
        num_beam_groups=num_beam_groups,
        diversity_penalty=diversity_penalty,
        temperature=temperature,
        top_k=top_k,
        top_p=top_p,
        return_dict_in_generate=True,
        output_scores=True,
    )

    sequences = outputs.sequences.reshape(
        -1, num_return_sequences, outputs.sequences.size(-1)
    )
    outputs.sequences = torch.zeros(
        sequences.size(0), num_return_sequences, max_gen_seq_length
    ).to(device=sequences.device)
    outputs.sequences[:, :, : sequences.size(-1)].copy_(sequences)

    if num_return_sequences == 1:
        outputs.sequences = outputs.sequences.reshape(-1, max_gen_seq_length)

    return GenericOutputs(
        sequences=outputs.sequences,
        sequences_scores=outputs.sequences_scores,
    )