Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def test_model_target_on_local(self):
obj = 1
file_path = os.path.join(_get_temporary_directory(), 'test.zip')
target = make_model_target(file_path=file_path,
temporary_directory=_get_temporary_directory(),
save_function=self._save_function,
load_function=self._load_function)
target.dump(obj)
loaded = target.load()
self.assertEqual(loaded, obj)
def test_model_target_on_s3(self):
conn = boto3.resource('s3', region_name='us-east-1')
conn.create_bucket(Bucket='test')
obj = 1
file_path = os.path.join('s3://test/', 'test.zip')
target = make_model_target(file_path=file_path,
temporary_directory=_get_temporary_directory(),
save_function=self._save_function,
load_function=self._load_function)
target.dump(obj)
loaded = target.load()
self.assertEqual(loaded, obj)
relative_file_path: str,
save_function: Callable[[Any, str], None],
load_function: Callable[[str], Any],
use_unique_id: bool = True):
"""
Make target for models which generate multiple files in saving, e.g. gensim.Word2Vec, Tensorflow, and so on.
:param relative_file_path: A file path to save.
:param save_function: A function to save a model. This takes a model object and a file path.
:param load_function: A function to load a model. This takes a file path and returns a model object.
:param use_unique_id: If this is true, add an unique id to a file base name.
"""
file_path = os.path.join(self.workspace_directory, relative_file_path)
assert relative_file_path[-3:] == 'zip', f'extension must be zip, but {relative_file_path} is passed.'
unique_id = self.make_unique_id() if use_unique_id else None
return gokart.target.make_model_target(file_path=file_path,
temporary_directory=self.local_temporary_directory,
unique_id=unique_id,
save_function=save_function,
load_function=load_function)
def make_large_data_frame_target(self, relative_file_path: str, use_unique_id: bool = True, max_byte=int(2**26)) -> TargetOnKart:
file_path = os.path.join(self.workspace_directory, relative_file_path)
unique_id = self.make_unique_id() if use_unique_id else None
return gokart.target.make_model_target(file_path=file_path,
temporary_directory=self.local_temporary_directory,
unique_id=unique_id,
save_function=gokart.target.LargeDataFrameProcessor(max_byte=max_byte).save,
load_function=gokart.target.LargeDataFrameProcessor.load)