unitorch.models.bart¤
BartProcessor¤
Bases: HfTextGenerationProcessor
Processor for BART model. Inherits from HfTextGenerationProcessor.
Initializes the BartProcessor.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
vocab_path |
str
|
Path to the vocabulary file. |
required |
merge_path |
str
|
Path to the BPE merge file. |
required |
special_input_ids |
Optional[Dict]
|
Optional dictionary of special input IDs. |
dict()
|
max_seq_length |
Optional[int]
|
Maximum sequence length. |
128
|
max_gen_seq_length |
Optional[int]
|
Maximum generation sequence length. |
48
|
Source code in src/unitorch/models/bart/processing.py
35 36 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 |
|
BartForGeneration¤
Bases: GenericModel
BART model for text generation.
Initializes a BartForGeneration model.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
config_path |
str
|
Path to the BART model configuration file. |
required |
gradient_checkpointing |
bool
|
Whether to use gradient checkpointing during training. Defaults to False. |
False
|
Source code in src/unitorch/models/bart/modeling.py
23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
|
forward ¤
forward(
input_ids: Tensor,
attention_mask: Tensor,
decoder_input_ids: Tensor,
decoder_attention_mask: Tensor,
)
Performs forward pass of the BartForGeneration model.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
input_ids |
Tensor
|
Tensor of input token IDs. |
required |
attention_mask |
Tensor
|
Tensor indicating which tokens should be attended to. |
required |
decoder_input_ids |
Tensor
|
Tensor of decoder input token IDs. |
required |
decoder_attention_mask |
Tensor
|
Tensor indicating which decoder tokens should be attended to. |
required |
Returns:
Type | Description |
---|---|
Tensor
|
Output logits of the model. |
Source code in src/unitorch/models/bart/modeling.py
42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
|
generate ¤
generate(
input_ids: Tensor,
num_beams: Optional[int] = 5,
decoder_start_token_id: Optional[int] = 2,
decoder_end_token_id: Optional[
Union[int, List[int]]
] = 2,
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 text using the BartForGeneration model.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
input_ids |
Tensor
|
Tensor of input token IDs. |
required |
num_beams |
int
|
Number of beams for beam search. Defaults to 5. |
5
|
decoder_start_token_id |
int
|
The ID of the decoder start token. Defaults to 2. |
2
|
decoder_end_token_id |
int or List[int]
|
The ID(s) of the decoder end token(s). Defaults to 2. |
2
|
num_return_sequences |
int
|
Number of sequences to return. Defaults to 1. |
1
|
min_gen_seq_length |
int
|
Minimum length of generated sequences. Defaults to 0. |
0
|
max_gen_seq_length |
int
|
Maximum length of generated sequences. Defaults to 48. |
48
|
repetition_penalty |
float
|
Repetition penalty for generated sequences. Defaults to 1.0. |
1.0
|
no_repeat_ngram_size |
int
|
Size of n-grams to avoid repetition. Defaults to 0. |
0
|
early_stopping |
bool
|
Whether to stop generation early. Defaults to True. |
True
|
length_penalty |
float
|
Length penalty for generated sequences. Defaults to 1.0. |
1.0
|
num_beam_groups |
int
|
Number of beam groups for diverse beam search. Defaults to 1. |
1
|
diversity_penalty |
float
|
Diversity penalty for diverse beam search. Defaults to 0.0. |
0.0
|
do_sample |
bool
|
Whether to use sampling during generation. Defaults to False. |
False
|
temperature |
float
|
Temperature for sampling. Defaults to 1.0. |
1.0
|
top_k |
int
|
Value for top-k sampling. Defaults to 50. |
50
|
top_p |
float
|
Value for top-p sampling. Defaults to 1.0. |
1.0
|
Returns:
Name | Type | Description |
---|---|---|
GenericOutputs |
Generated sequences and their scores. |
Source code in src/unitorch/models/bart/modeling.py
71 72 73 74 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 |
|