Spaces:
Runtime error
Runtime error
File size: 2,917 Bytes
f8f02c0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | 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
)
|