import torch
from transformers import GenerationConfig, PretrainedConfig
from transformers.generation.utils import GenerationMixin
from transformers.generation.continuous_batching.requests import GenerationOutput, RequestStatus

class DummyContinuousBatchingGenerateModel(GenerationMixin):
    def __init__(self):
        self.config = PretrainedConfig()
        self.generation_config = GenerationConfig()
        self.device = torch.device("cpu")

    def generate_batch(self, inputs, generation_config=None, **kwargs):
        num_return_sequences = generation_config.num_return_sequences or 1
        return {
            f"req_{i}": GenerationOutput(
                request_id=f"req_{i}",
                prompt_ids=inputs[0],
                generated_tokens=[10 + i],
                status=RequestStatus.FINISHED,
            )
            for i in range(num_return_sequences)
        }

model = DummyContinuousBatchingGenerateModel()
model.generate(
    inputs=torch.tensor([[1, 2, 3]]),
    cache_implementation="paged",
    do_sample=True,
    num_return_sequences=2,
)