@@ -42,7 +42,9 @@ class OpenlayerModel(abc.ABC):
4242 def run_from_cli (self ) -> None :
4343 """Run the model from the command line."""
4444 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+ )
4648 parser .add_argument (
4749 "--output-dir" ,
4850 type = str ,
@@ -85,7 +87,9 @@ def run_batch_from_df(self, df: pd.DataFrame) -> Tuple[pd.DataFrame, dict]:
8587 # Filter row_dict to only include keys that are valid parameters
8688 # for the 'run' method
8789 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+ }
8993
9094 # Call the run method with filtered kwargs
9195 output = self .run (** filtered_kwargs )
@@ -108,7 +112,8 @@ def run_batch_from_df(self, df: pd.DataFrame) -> Tuple[pd.DataFrame, dict]:
108112 if "tokens" in processed_trace :
109113 df .at [index , "tokens" ] = processed_trace ["tokens" ]
110114 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" ])
112117
113118 config = {
114119 "outputColumnName" : "output" ,
0 commit comments