Spaces:
Running
Running
| import os, re, io, time, json, logging, contextlib, textwrap | |
| from typing import Dict, Any | |
| import pandas as pd | |
| import argparse | |
| from auto_causal.agent import run_causal_analysis | |
| # Constants | |
| RATE_LIMIT_SECONDS = 2 | |
| def run_caia(desc, question, df): | |
| return run_causal_analysis(query=question, dataset_path=df, dataset_description=desc) | |
| def parse_args(): | |
| parser = argparse.ArgumentParser(description="Run batch causal analysis.") | |
| parser.add_argument("--csv_path", type=str, required=True, help="CSV file with queries, descriptions, and file names.") | |
| parser.add_argument("--data_folder", type=str, required=True, help="Folder containing data CSVs.") | |
| parser.add_argument("--data_category", type=str, required=True, help="Dataset category (e.g., real, qrdata, synthetic).") | |
| parser.add_argument("--output_folder", type=str, required=True, help="Folder to save output.") | |
| parser.add_argument("--llm_name", type=str, required=True, help="Name of the LLM used.") | |
| return parser.parse_args() | |
| def main(): | |
| args = parse_args() | |
| csv_meta = args.csv_meta | |
| data_dir = args.data_dir | |
| output_json = args.output_json | |
| os.environ["LLM_MODEL"] = args.llm_name | |
| print("[main] Starting batch processing…") | |
| if not os.path.exists(csv_meta): | |
| logging.error(f"Meta file not found: {csv_meta}") | |
| return | |
| meta_df = pd.read_csv(csv_meta) | |
| print(f"[main] Loaded metadata CSV with {len(meta_df)} rows.") | |
| results: Dict[int, Dict[str, Any]] = {} | |
| for idx, row in meta_df.iterrows(): | |
| data_path = os.path.join(data_dir, str(row["data_files"])) | |
| print(f"\n[main] Row {idx+1}/{len(meta_df)} → Dataset: {data_path}") | |
| try: | |
| res = run_caia( | |
| desc=row["data_description"], | |
| question=row["natural_language_query"], | |
| df=data_path, | |
| ) | |
| # Format result according to specified structure | |
| formatted_result = { | |
| "query": row["natural_language_query"], | |
| "method": row["method"], | |
| "answer": row["answer"], | |
| "dataset_description": row["data_description"], | |
| "dataset_path": data_path, | |
| "keywords": row.get("keywords", "Causality, Average treatment effect"), | |
| "final_result": { | |
| "method": res['results']['results'].get("method_used"), | |
| "causal_effect": res['results']['results'].get("effect_estimate"), | |
| "standard_deviation": res['results']['results'].get("standard_error"), | |
| "treatment_variable": res['results']['variables'].get("treatment_variable", None), | |
| "outcome_variable": res['results']['variables'].get("outcome_variable", None), | |
| "covariates": res['results']['variables'].get("covariates", []), | |
| "instrument_variable": res['results']['variables'].get("instrument_variable", None), | |
| "running_variable": res['results']['variables'].get("running_variable", None), | |
| "temporal_variable": res['results']['variables'].get("time_variable", None), | |
| "statistical_test_results": res.get("summary", ""), | |
| "explanation_for_model_choice": res.get("explanation", ""), | |
| "regression_equation": res.get("regression_equation", "") | |
| } | |
| } | |
| results[idx] = formatted_result | |
| print(type(res)) | |
| print(res) | |
| print(f"[main] Formatted result for row {idx+1}:", formatted_result) | |
| except Exception as e: | |
| logging.error(f"[{idx+1}] Error: {e}") | |
| results[idx] = {"answer": str(e)} | |
| time.sleep(RATE_LIMIT_SECONDS) | |
| os.makedirs(os.path.dirname(output_json), exist_ok=True) | |
| with open(output_json, "w") as f: | |
| json.dump(results, f, indent=2) | |
| print(f"[main] Done. Predictions saved to {output_json}") | |
| if __name__ == "__main__": | |
| main() | |