update

python/setup.py

@@ -4,9 +4,10 @@ setup(
     name="tabby",
     packages=find_packages(exclude=["tabby_tests"]),
     install_requires=[
+        "datasets",
         "dagster",
         "dagster-cloud",
-        "dagster-pandas"
+        "dagster-pandas",
     ],
     extras_require={"dev": ["dagster-webserver", "pytest"]},
 )

python/tabby/assets.py

@@ -24,7 +24,7 @@ DatasetDataFrame = create_dagster_pandas_dataframe_type(
 
 @asset(dagster_type=DatasetDataFrame)
 def dataset():
-    """Get source code information from TABBY_ROOT"""
+    """Read source code dataset from TABBY_ROOT"""
 
     ds = []
     for path in glob.glob(constants.TABBY_DATASET_FILEPATTERN):
@@ -50,34 +50,16 @@ def dataset():
     }
     return Output(df, metadata=metadata)
 
-EventDataFrame = create_dagster_pandas_dataframe_type(
+@asset
-    name="EventDataFrame",
+def train_dataset(dataset):
-    columns=[
+    """Filter source code dataset for training / evaluation"""
-        PandasColumn.integer_column("ts"),
+    from datasets import Dataset
-        PandasColumn.exists("event"),
-    ],
-)
 
-@asset(dagster_type=EventDataFrame)
+    df = dataset
-def events():
+    df = df[df["max_line_length"] < 300]
-    """Get events information from TABBY_ROOT"""
+    df = df[df["avg_line_length"] < 150]
-
-    ds = []
-    for path in glob.glob(constants.TABBY_EVENTS_FILEPATTERN):
-        with open(path, "r") as f:
-            for line in f.readlines():
-                ds.append(json.loads(line))
-
-    df = pd.DataFrame(ds)
     metadata = {
         "num_records": len(df),
-        "preview": MetadataValue.md(
+        "num_filtered_records": len(dataset) - len(df)
-            df.head()[
-                [
-                    "ts",
-                    "event"
-                ]
-            ].to_markdown()
-        ),
     }
-    return Output(df, metadata=metadata)
+    return Output(Dataset.from_pandas(df), metadata=metadata)