1  import json
 2  import os
 3  
 4  import pandas as pd
 5  from johnsnowlabs import nlp
 6  
 7  import mlflow
 8  from mlflow.pyfunc import spark_udf
 9  
10  # 1) Write your raw license.json string into the 'JOHNSNOWLABS_LICENSE_JSON' env variable for MLflow
11  creds = {
12      "AWS_ACCESS_KEY_ID": "...",
13      "AWS_SECRET_ACCESS_KEY": "...",
14      "SPARK_NLP_LICENSE": "...",
15      "SECRET": "...",
16  }
17  os.environ["JOHNSNOWLABS_LICENSE_JSON"] = json.dumps(creds)
18  
19  # 2) Install enterprise libraries
20  nlp.install()
21  # 3) Start a Spark session with enterprise libraries
22  spark = nlp.start()
23  
24  # 4) Load a model and test it
25  nlu_model = "en.classify.bert_sequence.covid_sentiment"
26  model_save_path = "my_model"
27  johnsnowlabs_model = nlp.load(nlu_model)
28  johnsnowlabs_model.predict(["I hate COVID,", "I love COVID"])
29  
30  # 5) Export model with pyfunc and johnsnowlabs flavors
31  with mlflow.start_run():
32      model_info = mlflow.johnsnowlabs.log_model(johnsnowlabs_model, name=model_save_path)
33  
34  # 6) Load model with johnsnowlabs flavor
35  mlflow.johnsnowlabs.load_model(model_info.model_uri)
36  
37  # 7) Load model with pyfunc flavor
38  mlflow.pyfunc.load_model(model_save_path)
39  
40  pandas_df = pd.DataFrame({"text": ["Hello World"]})
41  spark_df = spark.createDataFrame(pandas_df).coalesce(1)
42  pyfunc_udf = spark_udf(
43      spark=spark,
44      model_uri=model_save_path,
45      env_manager="virtualenv",
46      result_type="string",
47  )
48  new_df = spark_df.withColumn("prediction", pyfunc_udf(*pandas_df.columns))
49  
50  # 9) You can now use the mlflow models serve command to serve the model see next section
51  
52  # 10)  You can also use x command to deploy model inside of a container see next section