File size: 4,365 Bytes
6455f60 | 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 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | #!/usr/bin/env python3
"""
Data Streaming and Parsing Utility for the HackIndia Challenge.
Author: Team Ascended
Description: Memory-efficient streaming of financial QA datasets (specifically
mishface123/adaption-econ-finance-qa-pairs) with batch generation,
parsing, and formatting capabilities.
"""
import argparse
import sys
from typing import Dict, Generator, List, Optional
try:
from datasets import load_dataset
except ImportError:
print("Warning: 'datasets' library not found. Please install it via 'pip install datasets'.")
def stream_dataset(
dataset_name: str,
split: str = "train",
batch_size: int = 1,
limit: Optional[int] = None
) -> Generator[List[Dict], None, None]:
"""
Streams a dataset from Hugging Face and yields parsed data in batches.
Args:
dataset_name: Hugging Face dataset identifier.
split: The dataset split to stream (e.g., 'train', 'test').
batch_size: Number of records to yield at a time.
limit: Max number of records to process before stopping.
Yields:
A list of parsed dictionaries containing formatted instruction-response pairs.
"""
try:
# Load the dataset in streaming mode to minimize RAM overhead
dataset = load_dataset(dataset_name, split=split, streaming=True)
except Exception as e:
print(f"Error loading dataset {dataset_name}: {e}", file=sys.stderr)
return
batch = []
count = 0
for record in dataset:
parsed_record = parse_record(record)
batch.append(parsed_record)
count += 1
if len(batch) == batch_size:
yield batch
batch = []
if limit is not None and count >= limit:
break
if batch:
yield batch
def parse_record(record: Dict) -> Dict:
"""
Parses a single record from the dataset and standardizes its keys.
Designed specifically for 'mishface123/adaption-econ-finance-qa-pairs'.
Args:
record: The raw record dict from the dataset.
Returns:
A standardized dictionary containing instruction, context, and response.
"""
# Standardize keys depending on schema
instruction = record.get("instruction") or record.get("question") or ""
context = record.get("context") or ""
response = record.get("output") or record.get("response") or record.get("answer") or ""
metadata = record.get("metadata") or {}
# Extract target market metadata (e.g., Indian localization indicator)
localized = metadata.get("localized", False) or "india" in str(metadata).lower()
return {
"instruction": instruction.strip(),
"context": context.strip(),
"response": response.strip(),
"localized": localized,
"raw_keys": list(record.keys())
}
def main():
parser = argparse.ArgumentParser(
description="Stream and parse datasets for the AutoScientist pipeline."
)
parser.add_argument(
"--dataset",
type=str,
default="mishface123/adaption-econ-finance-qa-pairs",
help="Hugging Face dataset name to stream."
)
parser.add_argument(
"--split",
type=str,
default="train",
help="Dataset split (train, validation, test)."
)
parser.add_argument(
"--batch-size",
type=int,
default=4,
help="Batch size for yield/printing."
)
parser.add_argument(
"--limit",
type=int,
default=10,
help="Maximum records to stream/parse."
)
args = parser.parse_args()
print(f"Streaming dataset: {args.dataset} (split: {args.split})...")
batch_generator = stream_dataset(
dataset_name=args.dataset,
split=args.split,
batch_size=args.batch_size,
limit=args.limit
)
for i, batch in enumerate(batch_generator):
print(f"\n--- Batch {i+1} (Size: {len(batch)}) ---")
for record_idx, record in enumerate(batch):
print(f"\nRecord {record_idx+1}:")
print(f" Instruction : {record['instruction'][:120]}...")
print(f" Context : {record['context'][:120]}...")
print(f" Response : {record['response'][:120]}...")
print(f" Localized : {record['localized']}")
if __name__ == "__main__":
main()
|