SHARE
TWEET

Untitled

a guest Aug 25th, 2019 51 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. class PredictSklearn(beam.DoFn):
  2.     """ Format the input to the desired shape"""
  3.  
  4.     def __init__(self, project=None, bucket_name=None, model_path=None, destination_name=None):
  5.         self._model = None
  6.         self._project = project
  7.         self._bucket_name = bucket_name
  8.         self._model_path = model_path
  9.         self._destination_name = destination_name
  10.  
  11.     # Load once or very few times
  12.     def setup(self):
  13.         logging.info(
  14.             "Sklearn model initialisation {}".format(self._model_path))
  15.         download_blob(bucket_name=self._bucket_name, source_blob_name=self._model_path,
  16.                       project=self._project, destination_file_name=self._destination_name)
  17.         # unpickle sklearn model
  18.         self._model = pickle.load(open(self._destination_name, 'rb'))
  19.  
  20.     def process(self, element):
  21.         element["prediction"] = self._model.predict(element["data"])
  22.         return [element]
RAW Paste Data
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy. OK, I Understand
 
Top