Spaces:
Running on Zero
Running on Zero
File size: 4,195 Bytes
0418f40 | 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 | """Build the mixed finance instruction dataset in chat ("messages") format.
Each source dataset gets a formatter that maps its rows to
{"messages": [{"role": "user", ...}, {"role": "assistant", ...}], "source": name}
so TRL's SFTTrainer can apply any model's chat template at train time.
Sources that fail to load (renamed repo, script-based dataset on datasets>=3.0,
network) are skipped with a warning instead of killing the run.
Usage:
python -m src.data.prepare_datasets --out data/finance_sft \
[--max-per-source 20000] [--push-to-hub user/finance-sft-mix]
"""
import argparse
import random
from datasets import Dataset, concatenate_datasets, load_dataset
SYSTEM_PROMPT = (
"You are a financial analysis assistant with expertise in markets, filings, "
"accounting standards (IFRS/GAAP), and financial regulation. Be precise, cite "
"the relevant standard or filing section when applicable, and say so when you "
"are unsure. You do not give personalized investment advice."
)
def _chat(user, assistant, source):
return {
"messages": [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": user},
{"role": "assistant", "content": assistant},
],
"source": source,
}
def _fmt_instruction_style(row, source):
"""FinGPT-style rows: instruction / input / output."""
user = row["instruction"]
if row.get("input"):
user = f"{user}\n\n{row['input']}"
return _chat(user, row["output"], source)
def _fmt_finance_alpaca(row, source):
user = row["instruction"]
if row.get("input"):
user = f"{user}\n\n{row['input']}"
return _chat(user, row["output"], source)
def _fmt_phrasebank(row, source):
labels = {0: "negative", 1: "neutral", 2: "positive"}
user = (
"Classify the sentiment of this financial news sentence as positive, "
f"negative, or neutral:\n\n{row['sentence']}"
)
return _chat(user, labels[row["label"]], source)
# name -> (hf repo, config, split, formatter, mixture weight)
SOURCES = {
"fingpt-sentiment": ("FinGPT/fingpt-sentiment-train", None, "train", _fmt_instruction_style, 1.0),
"fingpt-fiqa-qa": ("FinGPT/fingpt-fiqa_qa", None, "train", _fmt_instruction_style, 1.5),
"fingpt-headline": ("FinGPT/fingpt-headline", None, "train", _fmt_instruction_style, 0.5),
"fingpt-finred": ("FinGPT/fingpt-finred", None, "train", _fmt_instruction_style, 0.5),
"finance-alpaca": ("gbharti/finance-alpaca", None, "train", _fmt_finance_alpaca, 1.5),
"financial-phrasebank": ("takala/financial_phrasebank", "sentences_66agree", "train", _fmt_phrasebank, 0.5),
}
def build(max_per_source, seed=42):
random.seed(seed)
parts = []
for name, (repo, config, split, fmt, weight) in SOURCES.items():
try:
ds = load_dataset(repo, config, split=split) if config else load_dataset(repo, split=split)
except Exception as e: # noqa: BLE001 - skip broken sources, keep the run alive
print(f"[skip] {name} ({repo}): {type(e).__name__}: {e}")
continue
n = min(len(ds), int(max_per_source * weight))
ds = ds.shuffle(seed=seed).select(range(n))
rows = [fmt(row, name) for row in ds]
parts.append(Dataset.from_list(rows))
print(f"[ok] {name}: {n} examples")
if not parts:
raise RuntimeError("No sources loaded — check network / dataset names.")
mixed = concatenate_datasets(parts).shuffle(seed=seed)
return mixed.train_test_split(test_size=0.02, seed=seed)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--out", default="data/finance_sft")
ap.add_argument("--max-per-source", type=int, default=20000)
ap.add_argument("--push-to-hub", default=None, help="e.g. user/finance-sft-mix (pushed private)")
args = ap.parse_args()
ds = build(args.max_per_source)
print(ds)
ds.save_to_disk(args.out)
print(f"Saved to {args.out}")
if args.push_to_hub:
ds.push_to_hub(args.push_to_hub, private=True)
print(f"Pushed to hub: {args.push_to_hub} (private)")
if __name__ == "__main__":
main()
|