From 55161f782c59f002557793a49cfa64cecc0240e3 Mon Sep 17 00:00:00 2001
From: JAX Toolbox <jax@nvidia.com>
Date: Sun, 20 Apr 2025 00:13:45 +0000
Subject: [PATCH] missing import

---
 pipe/etl.py | 163 ++++++++++++++++++++++++++--------------------------
 1 file changed, 83 insertions(+), 80 deletions(-)

diff --git a/pipe/etl.py b/pipe/etl.py
index a6be6c1..433e78e 100644
--- a/pipe/etl.py
+++ b/pipe/etl.py
@@ -1,80 +1,83 @@
-# -*- coding: utf-8 -*-
-"""
-etl.py
-
-This module contains the ETL (Extract, Transform, Load) pipeline for processing
-the spectrogram data and the labels.
-"""
-
-from pyspark.sql import SparkSession, functions, types, Row
-import tensorflow as tf
-import keras
-import matplotlib.pyplot as plt
-from sklearn.metrics import confusion_matrix
-from sklearn.preprocessing import OneHotEncoder
-
-def transform(spark, dataframe, keys):
-    dataframe = dataframe.withColumn( 
-            "index", functions.monotonically_increasing_id()
-    )
-    bundle = {key: [
-        arr.tolist()
-        for arr in OneHotEncoder(sparse_output=False) \
-            .fit_transform(dataframe.select(key).collect())
-        ] for key in keys
-    }
-
-    bundle = [dict(zip(bundle.keys(), values))
-              for values in zip(*bundle.values())]
-    schema = types.StructType([
-        types.StructField(key, types.ArrayType(types.FloatType()), True)
-        for key in keys
-    ])
-    newframe = spark.createDataFrame(bundle, schema=schema).withColumn(
-            "index", functions.monotonically_increasing_id()
-    )
-    for key in keys:
-        dataframe = dataframe.withColumnRenamed(key, f"{key}_str")
-    dataframe = dataframe.join(newframe, on="index", how="inner")
-
-    return dataframe
-
-def build_dict(df, key):
-    df = df.select(key, f"{key}_str").distinct()
-    
-    return df.rdd.map(
-        lambda row: (str(np.argmax(row[key])), row[f"{key}_str"])
-    ).collectAsMap()
-
-def extract(spark):
-    path = Path("/app/workdir")
-    labels = []
-    with open(path / "train.csv", "r") as file:
-        for line in file:
-            labels.append(line.strip().split(",")[0])
-
-    pipe = SpectrogramPipe(spark, filetype="matfiles")
-
-    return pipe.spectrogram_pipe(path, labels)
-
-def load(spark, split=[0.99, 0.005, 0.005]):
-    data = extract(spark)
-    data.select("treatment").replace("virus", "cpv") \
-                            .replace("cont", "pbs") \
-                            .replace("control", "pbs") \
-                            .replace("dld", "pbs").distinct()
-
-    data = transform(spark, data, ["treatment", "target"])
-    category_dict = {
-        key: build_dict(data, key) for key in ["treatment", "target"]
-    }
-    splits = data.randomSplit(split, seed=42)
-    trainx, valx, testx = (trim(dset, "spectrogram") for dset in splits)
-    trainy, valy, testy = (
-        [
-            np.array(dset.select("treatment").collect()).squeeze(),
-            np.array(dset.select("target").collect()).squeeze()
-        ] for dset in splits
-    )
-
-    return ((trainx, trainy), (valx, valy), (testx, testy), category_dict)
+# -*- coding: utf-8 -*-
+"""
+etl.py
+
+This module contains the ETL (Extract, Transform, Load) pipeline for processing
+the spectrogram data and the labels.
+"""
+
+import keras
+import matplotlib.pyplot as plt
+from pathlib import Path
+from pyspark.sql import SparkSession, functions, types, Row
+from sklearn.metrics import confusion_matrix
+from sklearn.preprocessing import OneHotEncoder
+import tensorflow as tf
+
+from pipe.pipe import SpectrogramPipe
+
+def transform(spark, dataframe, keys):
+    dataframe = dataframe.withColumn( 
+            "index", functions.monotonically_increasing_id()
+    )
+    bundle = {key: [
+        arr.tolist()
+        for arr in OneHotEncoder(sparse_output=False) \
+            .fit_transform(dataframe.select(key).collect())
+        ] for key in keys
+    }
+
+    bundle = [dict(zip(bundle.keys(), values))
+              for values in zip(*bundle.values())]
+    schema = types.StructType([
+        types.StructField(key, types.ArrayType(types.FloatType()), True)
+        for key in keys
+    ])
+    newframe = spark.createDataFrame(bundle, schema=schema).withColumn(
+            "index", functions.monotonically_increasing_id()
+    )
+    for key in keys:
+        dataframe = dataframe.withColumnRenamed(key, f"{key}_str")
+    dataframe = dataframe.join(newframe, on="index", how="inner")
+
+    return dataframe
+
+def build_dict(df, key):
+    df = df.select(key, f"{key}_str").distinct()
+    
+    return df.rdd.map(
+        lambda row: (str(np.argmax(row[key])), row[f"{key}_str"])
+    ).collectAsMap()
+
+def extract(spark):
+    path = Path("/app/workdir")
+    labels = []
+    with open(path / "train.csv", "r") as file:
+        for line in file:
+            labels.append(line.strip().split(",")[0])
+
+    pipe = SpectrogramPipe(spark, filetype="matfiles")
+
+    return pipe.spectrogram_pipe(path, labels)
+
+def load(spark, split=[0.99, 0.005, 0.005]):
+    data = extract(spark)
+    data.select("treatment").replace("virus", "cpv") \
+                            .replace("cont", "pbs") \
+                            .replace("control", "pbs") \
+                            .replace("dld", "pbs").distinct()
+
+    data = transform(spark, data, ["treatment", "target"])
+    category_dict = {
+        key: build_dict(data, key) for key in ["treatment", "target"]
+    }
+    splits = data.randomSplit(split, seed=42)
+    trainx, valx, testx = (trim(dset, "spectrogram") for dset in splits)
+    trainy, valy, testy = (
+        [
+            np.array(dset.select("treatment").collect()).squeeze(),
+            np.array(dset.select("target").collect()).squeeze()
+        ] for dset in splits
+    )
+
+    return ((trainx, trainy), (valx, valy), (testx, testy), category_dict)