@@ -42,7 +42,9 @@ class OpenlayerModel(abc.ABC):
42
42
def run_from_cli (self ) -> None :
43
43
"""Run the model from the command line."""
44
44
parser = argparse .ArgumentParser (description = "Run data through a model." )
45
- parser .add_argument ("--dataset-path" , type = str , required = True , help = "Path to the dataset" )
45
+ parser .add_argument (
46
+ "--dataset-path" , type = str , required = True , help = "Path to the dataset"
47
+ )
46
48
parser .add_argument (
47
49
"--output-dir" ,
48
50
type = str ,
@@ -85,7 +87,9 @@ def run_batch_from_df(self, df: pd.DataFrame) -> Tuple[pd.DataFrame, dict]:
85
87
# Filter row_dict to only include keys that are valid parameters
86
88
# for the 'run' method
87
89
row_dict = row .to_dict ()
88
- filtered_kwargs = {k : v for k , v in row_dict .items () if k in run_signature .parameters }
90
+ filtered_kwargs = {
91
+ k : v for k , v in row_dict .items () if k in run_signature .parameters
92
+ }
89
93
90
94
# Call the run method with filtered kwargs
91
95
output = self .run (** filtered_kwargs )
@@ -108,7 +112,8 @@ def run_batch_from_df(self, df: pd.DataFrame) -> Tuple[pd.DataFrame, dict]:
108
112
if "tokens" in processed_trace :
109
113
df .at [index , "tokens" ] = processed_trace ["tokens" ]
110
114
if "context" in processed_trace :
111
- df .at [index , "context" ] = processed_trace ["context" ]
115
+ # Convert the context list to a string to avoid pandas issues
116
+ df .at [index , "context" ] = json .dumps (processed_trace ["context" ])
112
117
113
118
config = {
114
119
"outputColumnName" : "output" ,
0 commit comments