train.py

"""
Financial Options & Market Prediction Expert Model
===================================================
Fine-tunes Mistral-7B-Instruct-v0.3 with QLoRA on ~745K financial instruction examples.
Combines 3 datasets:
  1. sujet-ai/Sujet-Finance-Instruct-177k (sentiment, NER, QA)
  2. gbharti/finance-alpaca (68K financial Q&A including options)
  3. Josephgflowers/Finance-Instruct-500k (500K broad financial instructions)

The model is trained with a system prompt emphasizing:
  - Options trading analysis
  - Explaining HOW data features affect market predictions
  - Step-by-step reasoning with feature importance
"""

import os
import torch
import trackio
from datasets import load_dataset, concatenate_datasets
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
)
from peft import LoraConfig, prepare_model_for_kbit_training
from trl import SFTTrainer, SFTConfig

# ============================================================================
# CONFIG
# ============================================================================
MODEL_ID = "mistralai/Mistral-7B-Instruct-v0.3"
HUB_MODEL_ID = "Saksham7772/FinOptions-Mistral-7B"
OUTPUT_DIR = "./finopt-mistral-7b-qlora"
PROJECT_NAME = "financial-options-expert"
RUN_NAME = "qlora-r64-lr2e4-ep2"

SYSTEM_PROMPT = (
    "You are a quantitative financial analyst and options trading expert. "
    "For every analysis you provide:\n"
    "1. Identify which input data features are most influential "
    "(e.g., implied volatility, volume, earnings, macro indicators, sentiment)\n"
    "2. Explain the directional impact of each feature on the prediction "
    "(bullish/bearish/neutral and why)\n"
    "3. Provide your market prediction or options strategy recommendation with clear reasoning\n"
    "4. Express your confidence level and key risk factors\n"
    "Think step by step before answering."
)

# ============================================================================
# TRACKIO — Experiment Tracking
# ============================================================================
trackio.init(
    project=PROJECT_NAME,
    name=RUN_NAME,
    config={
        "model": MODEL_ID,
        "lora_r": 64,
        "lora_alpha": 128,
        "learning_rate": 2e-4,
        "epochs": 2,
        "batch_size_per_device": 2,
        "gradient_accumulation_steps": 8,
        "effective_batch_size": 16,
        "quant": "4bit-nf4-double",
        "max_length": 2048,
        "datasets": [
            "sujet-ai/Sujet-Finance-Instruct-177k",
            "gbharti/finance-alpaca",
            "Josephgflowers/Finance-Instruct-500k",
        ],
    },
)

# ============================================================================
# DATASET PREPARATION
# ============================================================================
print("=" * 60)
print("Loading and converting datasets to messages format...")
print("=" * 60)

# --- Dataset 1: Sujet Finance 177K ---
ds_sujet = load_dataset("sujet-ai/Sujet-Finance-Instruct-177k", split="train")
print(f"  Sujet Finance: {len(ds_sujet)} rows")

def convert_sujet(example):
    system = example.get("system_prompt", "").strip()
    if not system:
        system = SYSTEM_PROMPT
    user = example.get("user_prompt", "").strip()
    answer = example.get("answer", "").strip()
    if not user or not answer:
        return {"messages": None}
    return {
        "messages": [
            {"role": "system", "content": system},
            {"role": "user", "content": user},
            {"role": "assistant", "content": answer},
        ]
    }

ds_sujet = ds_sujet.map(convert_sujet, remove_columns=ds_sujet.column_names, num_proc=4)
ds_sujet = ds_sujet.filter(lambda x: x["messages"] is not None, num_proc=4)
print(f"  Sujet Finance after conversion: {len(ds_sujet)} rows")

# --- Dataset 2: Finance Alpaca 68K ---
ds_alpaca = load_dataset("gbharti/finance-alpaca", split="train")
print(f"  Finance Alpaca: {len(ds_alpaca)} rows")

def convert_alpaca(example):
    instruction = example.get("instruction", "").strip()
    inp = example.get("input", "").strip()
    output = example.get("output", "").strip()
    if not instruction or not output:
        return {"messages": None}
    user_content = instruction
    if inp:
        user_content += f"\n\n{inp}"
    return {
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_content},
            {"role": "assistant", "content": output},
        ]
    }

ds_alpaca = ds_alpaca.map(convert_alpaca, remove_columns=ds_alpaca.column_names, num_proc=4)
ds_alpaca = ds_alpaca.filter(lambda x: x["messages"] is not None, num_proc=4)
print(f"  Finance Alpaca after conversion: {len(ds_alpaca)} rows")

# --- Dataset 3: Finance Instruct 500K ---
ds_500k = load_dataset("Josephgflowers/Finance-Instruct-500k", split="train")
print(f"  Finance Instruct 500K: {len(ds_500k)} rows")

def convert_500k(example):
    system = example.get("system", "").strip()
    if not system:
        system = SYSTEM_PROMPT
    user = example.get("user", "").strip()
    assistant = example.get("assistant", "").strip()
    if not user or not assistant:
        return {"messages": None}
    return {
        "messages": [
            {"role": "system", "content": system},
            {"role": "user", "content": user},
            {"role": "assistant", "content": assistant},
        ]
    }

ds_500k = ds_500k.map(convert_500k, remove_columns=ds_500k.column_names, num_proc=4)
ds_500k = ds_500k.filter(lambda x: x["messages"] is not None, num_proc=4)
print(f"  Finance Instruct 500K after conversion: {len(ds_500k)} rows")

# --- Combine all datasets ---
combined = concatenate_datasets([ds_sujet, ds_alpaca, ds_500k])
combined = combined.shuffle(seed=42)
print(f"\\n  COMBINED DATASET: {len(combined)} rows")

split = combined.train_test_split(test_size=0.01, seed=42)
train_dataset = split["train"]
eval_dataset = split["test"]
print(f"  Train: {len(train_dataset)} | Eval: {len(eval_dataset)}")
print("=" * 60)

# ============================================================================
# MODEL & TOKENIZER
# ============================================================================
print("Loading model with QLoRA 4-bit quantization...")

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    quantization_config=bnb_config,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
)
model = prepare_model_for_kbit_training(model)
model.config.use_cache = False

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

print(f"Model loaded: {MODEL_ID}")
print(f"Model params: {model.num_parameters():,} (quantized)")

# ============================================================================
# LoRA CONFIG — Following Open-FinLLMs recipe: r=64, alpha=128
# ============================================================================
peft_config = LoraConfig(
    r=64,
    lora_alpha=128,
    target_modules=[
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

# ============================================================================
# SFT CONFIG
# ============================================================================
sft_config = SFTConfig(
    output_dir=OUTPUT_DIR,
    num_train_epochs=2,
    per_device_train_batch_size=2,
    per_device_eval_batch_size=2,
    gradient_accumulation_steps=8,
    optim="paged_adamw_8bit",
    learning_rate=2e-4,
    max_grad_norm=0.3,
    weight_decay=0.001,
    warmup_ratio=0.03,
    lr_scheduler_type="cosine",
    bf16=True,
    fp16=False,
    max_length=2048,
    packing=False,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"use_reentrant": False},
    eval_strategy="steps",
    eval_steps=500,
    save_strategy="steps",
    save_steps=500,
    save_total_limit=3,
    load_best_model_at_end=True,
    metric_for_best_model="eval_loss",
    disable_tqdm=True,
    logging_strategy="steps",
    logging_steps=10,
    logging_first_step=True,
    logging_dir=f"{OUTPUT_DIR}/logs",
    report_to="trackio",
    run_name=RUN_NAME,
    push_to_hub=True,
    hub_model_id=HUB_MODEL_ID,
    hub_strategy="every_save",
    seed=42,
    dataloader_num_workers=4,
)

# ============================================================================
# TRAINER
# ============================================================================
trainer = SFTTrainer(
    model=model,
    args=sft_config,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    peft_config=peft_config,
    processing_class=tokenizer,
)

# ============================================================================
# TRAIN
# ============================================================================
print("\\n" + "=" * 60)
print("STARTING TRAINING")
print(f"  Model: {MODEL_ID}")
print(f"  Total train examples: {len(train_dataset)}")
print(f"  Epochs: 2")
print(f"  Effective batch size: 16")
print(f"  LoRA rank: 64, alpha: 128")
print(f"  Learning rate: 2e-4")
print(f"  Max sequence length: 2048")
print(f"  Push to Hub: {HUB_MODEL_ID}")
print("=" * 60 + "\\n")

trainer.train()

# ============================================================================
# SAVE & PUSH
# ============================================================================
print("\\nSaving final model...")
trainer.save_model(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)

print(f"\\nTraining complete! Model saved to Hub: https://huggingface.co/{HUB_MODEL_ID}")
trackio.finish()