""" 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()