import dill
from transformers import PreTrainedModel, PretrainedConfig

class CountAModel(PreTrainedModel):
    config_class = PretrainedConfig

    def __init__(self, config):
        super().__init__(config)

    def forward(self, text):
        return text.lower().count('a')

    def save_pretrained(self, save_directory):
        self.config.save_pretrained(save_directory)

config = PretrainedConfig()
config.torch_dtype = 'float32'  # Add a dummy torch_dtype attribute
config.model_type = 'CountA'
model = CountAModel(config)

# Validate
sentence = "This is a sample sentence with a few 'a's."
count_a = model(sentence)
print(f"The sentence contains {count_a} letter(s) 'a'.")

# Save the model in the current directory
model.save_pretrained(".")
with open('example-dummy-evaluation.dill', 'wb') as f:
    dill.dump(model, f)
    dill.dump(CountAModel, f)
print("Model saved successfully.")