Add serialization for DatasetInfo and round avg_words to int (#1294)

ludwig-ai · Sep 17, 2021 · 3ac82ea · 3ac82ea
1 parent 29cdc9b
commit 3ac82ea
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/ludwig/automl/base_config.py b/ludwig/automl/base_config.py
@@ -16,6 +16,7 @@
 
 import os
 from dataclasses import dataclass
+from dataclasses_json import LetterCase, dataclass_json
 from typing import List, Union
 
 import pandas as pd
@@ -45,6 +46,7 @@
 }
 
 
+@dataclass_json(letter_case=LetterCase.CAMEL)
 @dataclass
 class DatasetInfo:
     fields: List[FieldInfo]

diff --git a/ludwig/automl/utils.py b/ludwig/automl/utils.py
@@ -55,7 +55,7 @@ def avg_num_tokens(field: Series) -> int:
     if len(field) > 5000:
         field = field.sample(n=5000, random_state=40)
     unique_entries = field.unique()
-    avg_words = Series(unique_entries).str.split().str.len().mean()
+    avg_words = round(Series(unique_entries).str.split().str.len().mean())
     return avg_words