[T5] How to control the lenth of the generated summaries
Hi everyone, after fine-tuning a T5 model I wrote a script for inference:def generate_summary(text, max_length, temperature):
    inputs= tokenizer(text= text, return_tensors= "pt", truncation= True, padding= True, max_length= 1024).to(device)

    outputs= model.generate(
        **inputs,
        max_length= max_length,
        temperature= temperature,
        num_beams= 4,
        no_repeat_ngram_size= 2,
        early_stopping= True,
        do_sample= True
    )

    summary= tokenizer.decode(outputs[0], skip_special_tokens= True)

    return summaryHere in thegenerate_summarymethod, I parse three arguments:textfor input context,temperaturefor deterministic control, andmax_lengthfor output length control. However, I noticed that while thetemperatureworking well, themax_lengthwas not, it justcut off (truncated) the textwhen the maximum length was reached.Am I doing something wrong here?The goal is to tell the model to generate the output within the pre-set length.FYI, in the fine-tuning process, I created a Custom Dataset method:# LOAD AND PROCESS THE DATASET
class CustomDataset(Dataset):
    def __init__(self, data, context_max_length, summary_max_length):
        #self.data = data

        self.context_max_length= context_max_length
        self.summary_max_length= summary_max_length
        self.tokenizer= AutoTokenizer.from_pretrained(pretrained_model_name_or_path= parameters.model_name, cache_dir= parameters.model_cache_dir)
        
        self.data= [
            item for item in data if len(self.tokenizer.encode(item["context"]))<= self.context_max_length and len(self.tokenizer.encode(item["summary"]))<= self.summary_max_length
        ]

    def __len__(self):
        return len(self.data)  # Assuming both questions and passages have the same length
    
    def __getitem__(self, idx):
        context_encodings= self.tokenizer(
            self.data[idx]["context"],
            max_length= self.context_max_length,
            padding= "max_length",
            truncation= True,
            return_tensors= "pt"
        )

        summary_encodings= self.tokenizer(
            self.data[idx]["summary"],
            max_length= self.summary_max_length,
            padding= "max_length",
            truncation= True,
            return_tensors= "pt"
        )
        return {
            'input_ids': context_encodings["input_ids"].squeeze(), 
            'labels': summary_encodings["input_ids"].squeeze(),
            "attention_mask": context_encodings["attention_mask"].squeeze()
        }

print("Load the dataset")
train_dataset= load_from_disk(f"{parameters.datasets_cache_dir}/big_sum_dataset/train")
validation_dataset= load_from_disk(f"{parameters.datasets_cache_dir}/big_sum_dataset/validation")
tokenized_train_dataset= CustomDataset(data= train_dataset, context_max_length= 1024, summary_max_length= 512)
tokenized_validation_dataset= CustomDataset(data= validation_dataset, context_max_length= 1024, summary_max_length= 512)Here I do some preprocessing on the dataset:I filtered out all the instances that do not satisfy the pre-set lengthI padded the length of the input context to 1024 and the output summary to 512.I don’t know if this Custom Dataset makes the model can not generate vary length output.Thank you.