from transformers import PreTrainedModel, PretrainedConfig
import torch.nn as nn
import torch

class HelloWorldConfig(PretrainedConfig):
    model_type = "hello-world"

class HelloWorldModel(PreTrainedModel):
    config_class = HelloWorldConfig

    def __init__(self, config):
        super().__init__(config)
        # A single linear layer just for demo
        self.layer = nn.Linear(1, 1)

    def forward(self, input_ids=None):
        # Always returns “Hello World”
        return {"text": "Hello World"}

# Create the model + config
config = HelloWorldConfig()
model = HelloWorldModel(config)

# Save to local folder
model.save_pretrained("./")
config.save_pretrained("./")