|
| 1 | +# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +import argparse |
| 16 | +import json |
| 17 | +import zipfile |
| 18 | +from pathlib import Path |
| 19 | + |
| 20 | +from huggingface_hub import hf_hub_download |
| 21 | + |
| 22 | + |
| 23 | +def read_excel_to_text(excel_path: Path) -> str: |
| 24 | + """Read Excel file and convert to text representation.""" |
| 25 | + import pandas as pd |
| 26 | + |
| 27 | + try: |
| 28 | + # Explicitly handle .xlsb files with pyxlsb engine |
| 29 | + engine = "pyxlsb" if excel_path.suffix == ".xlsb" else None |
| 30 | + with pd.ExcelFile(excel_path, engine=engine) as xls: |
| 31 | + sheets = {sheet_name: xls.parse(sheet_name) for sheet_name in xls.sheet_names} |
| 32 | + except Exception as e: |
| 33 | + raise RuntimeError(f"Failed to read Excel file {excel_path}: {e}") from e |
| 34 | + |
| 35 | + combined_text = "" |
| 36 | + for sheet_name, df in sheets.items(): |
| 37 | + sheet_text = df.to_string(index=False) |
| 38 | + combined_text += f"Sheet name: {sheet_name}\n{sheet_text}\n\n" |
| 39 | + return combined_text |
| 40 | + |
| 41 | + |
| 42 | +def format_paths_for_prompt(paths: list[Path], actual_root: Path, display_root: Path) -> str: |
| 43 | + """Format file paths for display in prompt. |
| 44 | +
|
| 45 | + Args: |
| 46 | + paths: List of absolute Path objects to format |
| 47 | + actual_root: Root directory where files actually exist |
| 48 | + display_root: Root directory to display in paths (absolute for abs paths, Path(".") for relative) |
| 49 | + """ |
| 50 | + if not paths: |
| 51 | + return "" |
| 52 | + |
| 53 | + formatted = [] |
| 54 | + for path in paths: |
| 55 | + try: |
| 56 | + rel = path.relative_to(actual_root) |
| 57 | + disp_path = display_root / rel |
| 58 | + except ValueError: |
| 59 | + disp_path = path |
| 60 | + formatted.append(str(disp_path)) |
| 61 | + |
| 62 | + return " ".join(formatted) |
| 63 | + |
| 64 | + |
| 65 | +def save_data(split: str, data_dir: str | Path, display_root: str | Path | None, incontext_data: bool) -> None: |
| 66 | + """Download and prepare DSBench data.""" |
| 67 | + print(f"Preparing DSBench data for {split} split and saving to {data_dir}...") |
| 68 | + |
| 69 | + data_dir = Path(data_dir) |
| 70 | + data_dir.mkdir(parents=True, exist_ok=True) |
| 71 | + |
| 72 | + extracted_data_dir = data_dir / "data" |
| 73 | + |
| 74 | + # Extract if not already cached (hf_hub_download handles download caching) |
| 75 | + if not extracted_data_dir.exists(): |
| 76 | + print(" Downloading dataset from HuggingFace...") |
| 77 | + zip_path = Path( |
| 78 | + hf_hub_download(repo_id="liqiang888/DSBench", filename="data_analysis/data.zip", repo_type="dataset") |
| 79 | + ) |
| 80 | + print(" Extracting data...") |
| 81 | + with zipfile.ZipFile(zip_path, "r") as zip_ref: |
| 82 | + zip_ref.extractall(data_dir) |
| 83 | + if not extracted_data_dir.exists(): |
| 84 | + raise FileNotFoundError(f"Could not find data directory after extraction in {extracted_data_dir}") |
| 85 | + print(f" Dataset cached to {data_dir}") |
| 86 | + else: |
| 87 | + print(f" Using cached dataset from {data_dir}") |
| 88 | + |
| 89 | + # Load metadata |
| 90 | + print(" Loading metadata...") |
| 91 | + metadata_path = Path( |
| 92 | + hf_hub_download(repo_id="liqiang888/DSBench", filename="data_analysis/data.json", repo_type="dataset") |
| 93 | + ) |
| 94 | + metadata = [] |
| 95 | + with open(metadata_path, "r") as f: |
| 96 | + for line in f: |
| 97 | + if line.strip(): |
| 98 | + metadata.append(json.loads(line.strip())) |
| 99 | + |
| 100 | + # Process all tasks |
| 101 | + if not display_root: |
| 102 | + display_root = extracted_data_dir |
| 103 | + else: |
| 104 | + display_root = Path(display_root) |
| 105 | + |
| 106 | + print( |
| 107 | + f" Processing {len(metadata)} tasks at {extracted_data_dir} - using display root {display_root} for paths shown in the prompt..." |
| 108 | + ) |
| 109 | + all_entries = [] |
| 110 | + |
| 111 | + for task in metadata: |
| 112 | + task_id = task["id"] |
| 113 | + task_dir = extracted_data_dir / task_id |
| 114 | + |
| 115 | + if not task_dir.exists(): |
| 116 | + raise FileNotFoundError( |
| 117 | + f"Task directory not found: {task_dir}. " |
| 118 | + f"Expected task {task_id} from metadata but directory is missing. " |
| 119 | + "Data extraction may have failed." |
| 120 | + ) |
| 121 | + if len(task["answers"]) != len(task["questions"]): |
| 122 | + raise ValueError( |
| 123 | + f"Task {task_id}: mismatched questions ({len(task['questions'])}) " |
| 124 | + f"and answers ({len(task['answers'])}) counts in metadata." |
| 125 | + ) |
| 126 | + |
| 127 | + # Read introduction |
| 128 | + intro_file = task_dir / "introduction.txt" |
| 129 | + introduction = "" |
| 130 | + if intro_file.exists(): |
| 131 | + introduction = intro_file.read_text(encoding="utf-8", errors="ignore") |
| 132 | + |
| 133 | + # Get data files - support all Excel formats |
| 134 | + excel_files = [] |
| 135 | + for ext in ["*.xlsx", "*.xlsb", "*.xlsm"]: |
| 136 | + excel_files.extend(task_dir.glob(ext)) |
| 137 | + excel_files = [f for f in excel_files if "answer" not in f.name.lower()] |
| 138 | + |
| 139 | + # Read Excel content for in-context mode |
| 140 | + if incontext_data: |
| 141 | + excel_content = "" |
| 142 | + for excel_file in excel_files: |
| 143 | + sheets_text = read_excel_to_text(excel_file) |
| 144 | + excel_content += f"The excel file {excel_file.name} is: {sheets_text}\n\n" |
| 145 | + |
| 146 | + # Format paths for tool mode (relative to data directory) |
| 147 | + excel_paths = format_paths_for_prompt(excel_files, actual_root=extracted_data_dir, display_root=display_root) |
| 148 | + |
| 149 | + # Uncomment to get image files and csv files (for future multimodal and agentic support) |
| 150 | + # image_files = [] |
| 151 | + # for ext in ["*.jpg", "*.png", "*.jpeg"]: |
| 152 | + # image_files.extend(task_dir.glob(ext)) |
| 153 | + # csv_files = list(task_dir.glob("*.csv")) |
| 154 | + |
| 155 | + # Process each question |
| 156 | + for idx, question_name in enumerate(task["questions"]): |
| 157 | + question_file = task_dir / f"{question_name}.txt" |
| 158 | + |
| 159 | + if not question_file.exists(): |
| 160 | + print(f" Warning: {task_id}/{question_name}.txt not found, skipping") |
| 161 | + continue |
| 162 | + |
| 163 | + question_text = question_file.read_text(encoding="utf-8", errors="ignore").strip() |
| 164 | + |
| 165 | + # Build problem text (introduction + question) |
| 166 | + problem_text = "" |
| 167 | + if introduction: |
| 168 | + problem_text += f"The introduction is detailed as follows.\n{introduction}\n\n" |
| 169 | + problem_text += f"The question for this task is detailed as follows.\n{question_text}" |
| 170 | + |
| 171 | + # Create entry with all necessary fields |
| 172 | + entry = { |
| 173 | + # Skills standard fields |
| 174 | + "problem": problem_text, |
| 175 | + "expected_answer": task["answers"][idx], |
| 176 | + # For tool mode |
| 177 | + "excel_paths": excel_paths, |
| 178 | + # Metadata |
| 179 | + "task_id": task_id, |
| 180 | + "question_id": question_name, |
| 181 | + "task_name": task["name"], |
| 182 | + "task_url": task["url"], |
| 183 | + "task_year": task["year"], |
| 184 | + } |
| 185 | + |
| 186 | + if incontext_data: |
| 187 | + entry["excel_content"] = excel_content.strip() |
| 188 | + |
| 189 | + all_entries.append(entry) |
| 190 | + |
| 191 | + # Validate we got some entries |
| 192 | + if not all_entries: |
| 193 | + raise ValueError( |
| 194 | + f"No valid entries created! Processed {len(metadata)} tasks but all failed. " |
| 195 | + "Check that data was downloaded correctly and Excel files are readable." |
| 196 | + ) |
| 197 | + |
| 198 | + # Save to output file |
| 199 | + output_file = data_dir / f"{split}.jsonl" |
| 200 | + with open(output_file, "w") as f: |
| 201 | + for entry in all_entries: |
| 202 | + f.write(json.dumps(entry) + "\n") |
| 203 | + |
| 204 | + print(f" ✓ Saved {len(all_entries)} questions to {output_file}") |
| 205 | + |
| 206 | + |
| 207 | +if __name__ == "__main__": |
| 208 | + parser = argparse.ArgumentParser() |
| 209 | + parser.add_argument("--split", default="test", choices=("test",), help="DSBench only has test split") |
| 210 | + parser.add_argument( |
| 211 | + "--data_dir", type=str, default=None, help="Directory to save the data (defaults to dataset directory)" |
| 212 | + ) |
| 213 | + parser.add_argument( |
| 214 | + "--display_root", |
| 215 | + type=str, |
| 216 | + default=None, |
| 217 | + help='Root directory to display in paths (absolute for abs paths, Path(".") for relative)', |
| 218 | + ) |
| 219 | + parser.add_argument( |
| 220 | + "--incontext_data", |
| 221 | + action="store_true", |
| 222 | + help="Have the excel files read in-context under 'excel_content' field (Default: False)", |
| 223 | + ) |
| 224 | + args = parser.parse_args() |
| 225 | + print(args) |
| 226 | + if args.data_dir is None: |
| 227 | + # Save to the same directory as this script |
| 228 | + data_dir = Path(__file__).absolute().parent |
| 229 | + else: |
| 230 | + data_dir = Path(args.data_dir) |
| 231 | + |
| 232 | + save_data(args.split, data_dir, args.display_root, args.incontext_data) |
0 commit comments