macro_sentiment/finetune.py

"""
Fine-tune any sentiment head on your own labelled data.

Usage:
    python -m macro_sentiment.finetune \
        --data my_labels.csv \
        --text-column headline \
        --label-column sentiment \
        --base-model peyterho/finbert-macro-sentiment \
        --output peyterho/my-custom-model \
        --epochs 4

Accepts CSV, TSV, JSON, or JSONL files.
Labels can be strings ("positive", "negative", "neutral") or integers (0, 1, 2).

Label conventions:
    0 or "negative"  → negative sentiment
    1 or "neutral"   → neutral sentiment
    2 or "positive"  → positive sentiment
"""

import os
import sys
import json
import argparse
import numpy as np
import torch
from collections import Counter
from pathlib import Path

from datasets import load_dataset, Dataset, DatasetDict, ClassLabel
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    TrainingArguments,
    Trainer,
    DataCollatorWithPadding,
)
import evaluate


# ─── Label handling ───────────────────────────────────────────────

MODEL_LABEL_MAPS = {
    "ProsusAI/finbert": {0: 1, 1: 2, 2: 0},
    "peyterho/finbert-macro-sentiment": {0: 1, 1: 2, 2: 0},
    "soleimanian/financial-roberta-large-sentiment": None,
    "peyterho/financial-roberta-large-macro-sentiment": None,
    "climatebert/distilroberta-base-climate-sentiment": {0: 2, 1: 1, 2: 0},
    "peyterho/climatebert-macro-sentiment": {0: 2, 1: 1, 2: 0},
}

STRING_LABEL_MAP = {
    "negative": 0, "neg": 0, "bearish": 0, "risk": 0,
    "neutral": 1, "neut": 1, "mixed": 1,
    "positive": 2, "pos": 2, "bullish": 2, "opportunity": 2,
}


def parse_label(label):
    """Convert string or int label to unified int (0=neg, 1=neutral, 2=pos)."""
    if isinstance(label, (int, float)):
        label = int(label)
        if label in (0, 1, 2):
            return label
        raise ValueError(f"Integer label must be 0, 1, or 2. Got: {label}")
    label_lower = str(label).strip().lower()
    if label_lower in STRING_LABEL_MAP:
        return STRING_LABEL_MAP[label_lower]
    raise ValueError(
        f"Unknown label: '{label}'. "
        f"Accepted: {list(STRING_LABEL_MAP.keys())} or integers 0/1/2."
    )


# ─── Data loading ─────────────────────────────────────────────────

def load_user_data(data_path, text_column, label_column, val_split=0.15, seed=42):
    """Load user data from CSV/TSV/JSON/JSONL file."""
    path = Path(data_path)
    ext = path.suffix.lower()
    if ext == ".csv":
        ds = load_dataset("csv", data_files=str(path), split="train")
    elif ext == ".tsv":
        ds = load_dataset("csv", data_files=str(path), delimiter="\t", split="train")
    elif ext in (".json", ".jsonl"):
        ds = load_dataset("json", data_files=str(path), split="train")
    else:
        raise ValueError(f"Unsupported format: {ext}. Use .csv, .tsv, .json, or .jsonl")

    if text_column not in ds.column_names:
        raise ValueError(f"Text column '{text_column}' not found. Available: {ds.column_names}")
    if label_column not in ds.column_names:
        raise ValueError(f"Label column '{label_column}' not found. Available: {ds.column_names}")

    texts, labels, errors = [], [], []
    for i, row in enumerate(ds):
        text = str(row[text_column]).strip()
        if not text or len(text) < 3:
            continue
        try:
            label = parse_label(row[label_column])
            texts.append(text)
            labels.append(label)
        except ValueError as e:
            errors.append(f"  Row {i}: {e}")
            if len(errors) > 10:
                break

    if errors:
        print(f"\n⚠️  Label errors ({len(errors)} rows):")
        for e in errors[:11]:
            print(e)
    if not texts:
        raise ValueError("No valid samples found. Check column names and label format.")

    full_ds = Dataset.from_dict({"text": texts, "label": labels})
    full_ds = full_ds.cast_column("label", ClassLabel(names=["negative", "neutral", "positive"]))
    split = full_ds.train_test_split(test_size=val_split, seed=seed, stratify_by_column="label")
    return DatasetDict({"train": split["train"], "test": split["test"]})


# ─── Weighted Trainer ─────────────────────────────────────────────

class WeightedTrainer(Trainer):
    def __init__(self, class_weights=None, **kwargs):
        super().__init__(**kwargs)
        self.class_weights = class_weights

    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        labels = inputs.pop("labels")
        outputs = model(**inputs)
        logits = outputs.logits
        if self.class_weights is not None:
            loss_fct = torch.nn.CrossEntropyLoss(weight=self.class_weights.to(logits.device))
        else:
            loss_fct = torch.nn.CrossEntropyLoss()
        loss = loss_fct(logits, labels)
        return (loss, outputs) if return_outputs else loss


# ─── Auto-detect label remap ─────────────────────────────────────

def detect_label_remap(model, base_model_name):
    """Determine label remapping from model config."""
    if base_model_name in MODEL_LABEL_MAPS:
        return MODEL_LABEL_MAPS[base_model_name]

    config = model.config
    if not hasattr(config, "id2label"):
        return None

    id2label = {int(k): v.lower() for k, v in config.id2label.items()}
    if id2label.get(0) in ("negative", "neg") and id2label.get(2) in ("positive", "pos"):
        return None  # matches unified

    model_label_to_idx = {}
    for idx, name in id2label.items():
        if name in ("positive", "pos", "opportunity"):
            model_label_to_idx["pos"] = idx
        elif name in ("negative", "neg", "risk"):
            model_label_to_idx["neg"] = idx
        elif name in ("neutral", "neut"):
            model_label_to_idx["neut"] = idx

    if len(model_label_to_idx) == 3:
        return {0: model_label_to_idx["neg"], 1: model_label_to_idx["neut"], 2: model_label_to_idx["pos"]}

    print(f"⚠️  Could not auto-detect label mapping from: {config.id2label}")
    return None


# ─── Main ─────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(
        description="Fine-tune a sentiment model on your own labelled data.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python -m macro_sentiment.finetune \\
      --data my_labels.csv --text-column text --label-column sentiment

  python -m macro_sentiment.finetune \\
      --data data.jsonl --text-column body --label-column label \\
      --base-model peyterho/financial-roberta-large-macro-sentiment \\
      --output my-org/my-model --push-to-hub

  python -m macro_sentiment.finetune \\
      --data labels.csv --text-column headline --label-column tone \\
      --epochs 8 --lr 1e-5 --batch-size 16 --max-length 256
        """,
    )
    parser.add_argument("--data", required=True, help="Path to labelled data (.csv/.tsv/.json/.jsonl)")
    parser.add_argument("--text-column", required=True, help="Column containing text")
    parser.add_argument("--label-column", required=True, help="Column containing labels")
    parser.add_argument("--base-model", default="peyterho/finbert-macro-sentiment",
                        help="Base model (default: peyterho/finbert-macro-sentiment)")
    parser.add_argument("--output", default="./custom-finetuned", help="Output dir or Hub ID")
    parser.add_argument("--push-to-hub", action="store_true", help="Push to HF Hub")
    parser.add_argument("--epochs", type=int, default=4)
    parser.add_argument("--lr", type=float, default=2e-5)
    parser.add_argument("--batch-size", type=int, default=32)
    parser.add_argument("--max-length", type=int, default=128)
    parser.add_argument("--val-split", type=float, default=0.15)
    parser.add_argument("--no-class-weights", action="store_true")
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    print("=" * 60)
    print("Custom Fine-Tuning")
    print("=" * 60)
    print(f"  Data:       {args.data}")
    print(f"  Text col:   {args.text_column}")
    print(f"  Label col:  {args.label_column}")
    print(f"  Base model: {args.base_model}")
    print(f"  Output:     {args.output}")
    print(f"  Epochs:     {args.epochs}, LR: {args.lr}, Batch: {args.batch_size}")

    # Load data
    ds = load_user_data(args.data, args.text_column, args.label_column, args.val_split, args.seed)
    train_dist = Counter(ds["train"]["label"])
    print(f"\n  Train: {len(ds['train'])} samples — {dict(train_dist)}")
    print(f"  Val:   {len(ds['test'])} samples")

    # Load model
    tokenizer = AutoTokenizer.from_pretrained(args.base_model)
    model = AutoModelForSequenceClassification.from_pretrained(args.base_model, num_labels=3)
    label_remap = detect_label_remap(model, args.base_model)
    print(f"  Label remap: {label_remap or 'none needed'}")

    # Tokenize
    def preprocess(examples):
        tok = tokenizer(examples["text"], truncation=True, max_length=args.max_length)
        tok["labels"] = [label_remap[l] for l in examples["label"]] if label_remap else examples["label"]
        return tok

    tokenized = ds.map(preprocess, batched=True, remove_columns=["text", "label"])

    # Class weights
    class_weights = None
    if not args.no_class_weights:
        lc = Counter(tokenized["train"]["labels"])
        nt = sum(lc.values())
        class_weights = torch.tensor([np.sqrt(nt / lc.get(i, 1)) for i in range(3)], dtype=torch.float32)
        class_weights = class_weights / class_weights.mean()
        print(f"  Class weights: [{', '.join(f'{w:.3f}' for w in class_weights.tolist())}]")

    # Metrics
    acc_m = evaluate.load("accuracy")
    f1_m = evaluate.load("f1")

    def compute_metrics(ep):
        logits, labels = ep
        preds = np.argmax(logits, axis=-1)
        return {
            "accuracy": acc_m.compute(predictions=preds, references=labels)["accuracy"],
            "f1_macro": f1_m.compute(predictions=preds, references=labels, average="macro")["f1"],
            "f1_weighted": f1_m.compute(predictions=preds, references=labels, average="weighted")["f1"],
        }

    is_hub = args.push_to_hub or ("/" in args.output and not args.output.startswith("./"))
    output_dir = args.output if not is_hub else f"./finetune-{args.output.replace('/', '-')}"

    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=args.epochs,
        per_device_train_batch_size=args.batch_size,
        per_device_eval_batch_size=min(64, args.batch_size * 2),
        gradient_accumulation_steps=max(1, 64 // args.batch_size),
        learning_rate=args.lr,
        weight_decay=0.01,
        warmup_ratio=0.2,
        lr_scheduler_type="linear",
        eval_strategy="epoch",
        save_strategy="epoch",
        load_best_model_at_end=True,
        metric_for_best_model="f1_macro",
        greater_is_better=True,
        bf16=torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8,
        fp16=torch.cuda.is_available() and torch.cuda.get_device_capability()[0] < 8,
        logging_strategy="steps",
        logging_steps=max(1, len(tokenized["train"]) // (args.batch_size * 10)),
        logging_first_step=True,
        push_to_hub=is_hub,
        hub_model_id=args.output if is_hub else None,
        save_total_limit=2,
        report_to="none",
        seed=args.seed,
    )

    trainer = WeightedTrainer(
        class_weights=class_weights,
        model=model,
        args=training_args,
        train_dataset=tokenized["train"],
        eval_dataset=tokenized["test"],
        processing_class=tokenizer,
        data_collator=DataCollatorWithPadding(tokenizer=tokenizer),
        compute_metrics=compute_metrics,
    )

    print(f"\nTraining for {args.epochs} epochs...")
    trainer.train()

    ev = trainer.evaluate()
    print(f"\nResults:  Accuracy={ev['eval_accuracy']:.4f}  F1 macro={ev['eval_f1_macro']:.4f}  F1 weighted={ev['eval_f1_weighted']:.4f}")

    if is_hub:
        trainer.push_to_hub(commit_message=f"Fine-tuned {args.base_model} on custom data")
        print(f"✅ Model at: https://huggingface.co/{args.output}")
    else:
        trainer.save_model(output_dir)
        tokenizer.save_pretrained(output_dir)
        print(f"✅ Saved to: {output_dir}")

    with open(os.path.join(output_dir, "eval_results.json"), "w") as f:
        json.dump(ev, f, indent=2)


if __name__ == "__main__":
    main()