Rename 'predictions' column to 'output' for all model runners

Author: gustavocidornelas

examples/llms/general-llm/product-names.ipynb

@@ -356,7 +356,7 @@
    "outputs": [],
    "source": [
     "# There are costs in running this cell!\n",
-    "dataset[\"model_output\"] = llm_runner.run(dataset)[\"predictions\"]"
+    "dataset[\"model_output\"] = llm_runner.run(dataset)[\"output\"]"
    ]
   },
   {

examples/llms/ner/entity-extraction.ipynb

@@ -389,7 +389,7 @@
    "outputs": [],
    "source": [
     "# There are costs in running this cell!\n",
-    "dataset[\"model_output\"] = llm_runner.run(dataset)[\"predictions\"]"
+    "dataset[\"model_output\"] = llm_runner.run(dataset)[\"output\"]"
    ]
   },
   {

openlayer/model_runners/ll_model_runners.py

@@ -104,7 +104,7 @@ def _run_in_memory_and_yield_progress(
                     model_outputs.append(output_value)
                     current_row += 1
                     yield pd.DataFrame(
-                        {"predictions": model_outputs, "output_time_utc": timestamps}
+                        {"output": model_outputs, "output_time_utc": timestamps}
                     ), current_row / total_rows
                     continue
 
@@ -117,7 +117,7 @@ def _run_in_memory_and_yield_progress(
             current_row += 1
 
             yield pd.DataFrame(
-                {"predictions": model_outputs, "output_time_utc": timestamps}
+                {"output": model_outputs, "output_time_utc": timestamps}
             ), current_row / total_rows
 
         self.logger.info("Successfully ran data through the model!")
@@ -126,7 +126,7 @@ def _run_in_memory_and_yield_progress(
         self.cost_estimates.append(run_cost)
 
         yield pd.DataFrame(
-            {"predictions": model_outputs, "output_time_utc": timestamps}
+            {"output": model_outputs, "output_time_utc": timestamps}
         ), 1.0
 
     def _run_single_input(self, input_data_row: pd.Series) -> Tuple[str, float, set]:

openlayer/model_runners/prediction_jobs/classification_prediction_job.py

@@ -34,9 +34,7 @@
 
     # Run model
     logger.debug("Running model...")
-    output_data = pd.DataFrame(
-        {"predictions": ml_model.predict_proba(input_data).tolist()}
-    )
+    output_data = pd.DataFrame({"output": ml_model.predict_proba(input_data).tolist()})
 
     # Save output data
     logger.debug("Saving output data...")

openlayer/model_runners/prediction_jobs/regression_prediction_job.py

@@ -34,7 +34,7 @@
 
     # Run model
     logger.debug("Running model...")
-    output_data = pd.DataFrame({"predictions": ml_model.predict(input_data)})
+    output_data = pd.DataFrame({"output": ml_model.predict(input_data)})
 
     # Save output data
     logger.debug("Saving output data...")

openlayer/model_runners/traditional_ml_model_runners.py

@@ -110,9 +110,9 @@ def _post_process_output(self, output_data: pd.DataFrame) -> pd.DataFrame:
         processed_output_data = output_data.copy()
 
         # Make the items list of floats (and not strings)
-        processed_output_data["predictions"] = processed_output_data[
-            "predictions"
-        ].apply(ast.literal_eval)
+        processed_output_data["output"] = processed_output_data["output"].apply(
+            ast.literal_eval
+        )
 
         return processed_output_data