Sync SAL sources and dependencies

GitOrigin-RevId: dea281a81071fed5dab34b66b0cbc861e6cb3911

Author: mikeknep

notebooks/conditional-generation.py

@@ -1,24 +1,29 @@
 import pandas as pd
-from gretel_client import configure_session
 
+from gretel_client import configure_session
 from gretel_trainer import Trainer
-from gretel_trainer.models import GretelLSTM, GretelACTGAN
+from gretel_trainer.models import GretelACTGAN, GretelLSTM
 
-DATASET_PATH = 'https://gretel-public-website.s3.amazonaws.com/datasets/mitre-synthea-health.csv'
+DATASET_PATH = (
+    "https://gretel-public-website.s3.amazonaws.com/datasets/mitre-synthea-health.csv"
+)
 MODEL_TYPE = [GretelLSTM(), GretelACTGAN()][1]
 
 # Create dataset to autocomplete values for
-seed_df = pd.DataFrame(data=[
-    ["black", "african", "F"],
-    ["black", "african", "F"],
-    ["black", "african", "F"],
-    ["black", "african", "F"],
-    ["asian", "chinese", "F"],
-    ["asian", "chinese", "F"],
-    ["asian", "chinese", "F"],
-    ["asian", "chinese", "F"],
-    ["asian", "chinese", "F"]
-], columns=["RACE", "ETHNICITY", "GENDER"])
+seed_df = pd.DataFrame(
+    data=[
+        ["black", "african", "F"],
+        ["black", "african", "F"],
+        ["black", "african", "F"],
+        ["black", "african", "F"],
+        ["asian", "chinese", "F"],
+        ["asian", "chinese", "F"],
+        ["asian", "chinese", "F"],
+        ["asian", "chinese", "F"],
+        ["asian", "chinese", "F"],
+    ],
+    columns=["RACE", "ETHNICITY", "GENDER"],
+)
 
 
 # Configure Gretel credentials
@@ -31,5 +36,5 @@
 print(model.generate(seed_df=seed_df))
 
 # Load a existing model and conditionally generate data
-#model = Trainer.load()
-#print(model.generate(seed_df=seed_df))
+# model = Trainer.load()
+# print(model.generate(seed_df=seed_df))

notebooks/custom-example.py

@@ -9,10 +9,7 @@
 
 # Specify underlying model and config options.
 # configs can be either a string, dict, or path
-model_type = GretelACTGAN(
-    config="synthetics/tabular-actgan",
-    max_rows=50000
-)
+model_type = GretelACTGAN(config="synthetics/tabular-actgan", max_rows=50000)
 
 # Optionally update model parameters from a base config
 model_type.update_params({"epochs": 500})

notebooks/simple-example.py

@@ -13,5 +13,5 @@
 
 # Or, load and generate data from an existing model
 
-#model = Trainer.load()
-#model.generate(num_records=70)
+# model = Trainer.load()
+# model.generate(num_records=70)

requirements.txt

@@ -9,7 +9,7 @@ plotly~=5.11
 pydantic~=1.9
 requests~=2.25
 scikit-learn~=1.0
-smart-open[s3]~=5.2
+smart_open[s3]~=5.2
 sqlalchemy~=1.4
-typing-extensions~=4.7
+typing_extensions~=4.7
 unflatten==0.1.1

setup.py

@@ -1,5 +1,6 @@
 import pathlib
-from setuptools import setup, find_packages
+
+from setuptools import find_packages, setup
 
 local_path = pathlib.Path(__file__).parent
 install_requires = (local_path / "requirements.txt").read_text().splitlines()

src/gretel_trainer/benchmark/core.py

@@ -1,6 +1,7 @@
 import csv
 import logging
 import time
+
 from dataclasses import dataclass, field
 from datetime import datetime
 from enum import Enum

src/gretel_trainer/benchmark/custom/datasets.py

@@ -1,6 +1,7 @@
 import logging
 import os
 import uuid
+
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Optional, Union

src/gretel_trainer/benchmark/custom/strategy.py

@@ -1,7 +1,7 @@
 from pathlib import Path
 from typing import Optional
 
-from gretel_trainer.benchmark.core import BenchmarkConfig, Dataset, Timer, run_out_path
+from gretel_trainer.benchmark.core import BenchmarkConfig, Dataset, run_out_path, Timer
 from gretel_trainer.benchmark.custom.models import CustomModel
 
 

src/gretel_trainer/benchmark/entrypoints.py

@@ -2,21 +2,22 @@
 
 import logging
 import shutil
+
 from inspect import isclass
 from pathlib import Path
-from typing import Optional, Type, Union, cast
+from typing import cast, Optional, Type, Union
 
 import pandas as pd
-from gretel_client.config import get_session_config
 
+from gretel_client.config import get_session_config
 from gretel_trainer.benchmark.core import BenchmarkConfig, BenchmarkException, Dataset
 from gretel_trainer.benchmark.custom.models import CustomModel
 from gretel_trainer.benchmark.gretel.models import GretelModel
 from gretel_trainer.benchmark.job_spec import (
     DatasetTypes,
     JobSpec,
-    ModelTypes,
     model_name,
+    ModelTypes,
 )
 from gretel_trainer.benchmark.session import Session
 

src/gretel_trainer/benchmark/executor.py

@@ -1,10 +1,10 @@
 import logging
+
 from enum import Enum
 from typing import Optional, Protocol
 
 from gretel_client.projects.models import Model
 from gretel_client.projects.projects import Project
-
 from gretel_trainer.benchmark.core import BenchmarkConfig, Dataset, log, run_out_path
 from gretel_trainer.benchmark.sdk_extras import create_evaluate_model, run_evaluate