from transformers import T5Tokenizer, T5ForConditionalGeneration

class GenerationModel:
    def __init__(self, model_name):
        self.tokenizer = T5Tokenizer.from_pretrained(model_name)
        self.model = T5ForConditionalGeneration.from_pretrained(model_name)

    def generate(self, data):
        inputs = self.tokenizer.encode_plus(
            data,
            add_special_tokens=True,
            max_length=512,
            return_attention_mask=True,
            return_tensors="pt"
        )

        outputs = self.model.generate(inputs["input_ids"], attention_mask=inputs["attention_mask"])
        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)