flavienbwk

Now using dill file

59bf28c over 2 years ago

898 Bytes

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.")

	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.")