Spaces:
Runtime error
Runtime error
| from typing import Any, Dict | |
| import pandas as pd | |
| from .schema import TransformedOutput | |
| def parse_model_response( | |
| original_df: pd.DataFrame, | |
| api_response: Dict[str, Any] | |
| ) -> TransformedOutput: | |
| """ | |
| Parses the completed API response and aligns it strictly | |
| with the original input data. | |
| """ | |
| # 1. Extract predictions block from the API response | |
| api_preds = [] | |
| # The results are typically nested under "data" | |
| data_block = api_response.get("data", api_response) | |
| if "predictions" in data_block: | |
| outer_preds = data_block["predictions"] | |
| if isinstance(outer_preds, dict) and "predictions" in outer_preds: | |
| api_preds = outer_preds["predictions"] | |
| elif isinstance(outer_preds, list): | |
| api_preds = outer_preds | |
| if not api_preds: | |
| raise ValueError("No predictions returned in API response payload") | |
| if len(api_preds) != len(original_df): | |
| raise ValueError( | |
| f"Row mismatch: Input has {len(original_df)} rows, " | |
| f"but predictions have {len(api_preds)} rows." | |
| ) | |
| pred_df = pd.DataFrame(api_preds) | |
| # 2. Strict row alignment logic | |
| # If a specific merge column 'id' is present in both, merge by ID to guarantee order and alignment | |
| # Otherwise, fallback to row-index concatenation. | |
| merge_col = "id" | |
| if merge_col in original_df.columns and merge_col in pred_df.columns: | |
| enriched_df = pd.merge(original_df, pred_df, on=merge_col, how="left") | |
| else: | |
| # Strict row alignment without ID: ensure lengths match | |
| if len(original_df) != len(pred_df): | |
| raise ValueError( | |
| f"Row mismatch: Input has {len(original_df)} rows, " | |
| f"but predictions have {len(pred_df)} rows." | |
| ) | |
| # Reset index to guarantee correct horizontal concat | |
| df_left = original_df.reset_index(drop=True) | |
| df_right = pred_df.reset_index(drop=True) | |
| enriched_df = pd.concat([df_left, df_right], axis=1) | |
| # Convert all NaN values to None for clean JSON serialization | |
| input_data = original_df.where(pd.notnull(original_df), None).to_dict(orient="records") | |
| # Add row_index to predictions if not present, to fulfill structured output req | |
| structured_predictions = [] | |
| for idx, row in pred_df.iterrows(): | |
| structured_predictions.append({ | |
| "row_index": idx, | |
| "predicted_label": row.get("predicted_label") or row.get("label"), | |
| "probability": row.get("probability") or row.get("score") | |
| }) | |
| enriched_output = enriched_df.where(pd.notnull(enriched_df), None).to_dict(orient="records") | |
| return TransformedOutput( | |
| input_data=input_data, | |
| predictions=structured_predictions, | |
| enriched_output=enriched_output, | |
| enriched_df=enriched_df | |
| ) | |