|
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' |
|
config.model_type = 'CountA' |
|
model = CountAModel(config) |
|
|
|
|
|
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'.") |
|
|
|
|
|
model.save_pretrained(".") |
|
with open('example-dummy-evaluation.dill', 'wb') as f: |
|
dill.dump(model, f) |
|
dill.dump(CountAModel, f) |
|
print("Model saved successfully.") |
|
|
