Add custom model_name support to all heads + update defaults to fine-tuned models

macro_sentiment/transformers_ensemble.py

@@ -2,12 +2,17 @@
 Multi-head transformer ensemble for macroeconomic sentiment.
 
 Domain-specific heads:
-  1. FinBERT (ProsusAI) — financial news sentiment
+  1. FinBERT — financial news sentiment
   2. Financial-RoBERTa-large — policy/formal text sentiment
   3. ClimateBERT — climate risk/opportunity
   4. Keyword-based Topic Router — routes to appropriate head
 
 All heads output standardized scores in [-1, +1].
+
+Fine-tuned models (v0.2.0):
+  - FinBERT:    peyterho/finbert-macro-sentiment
+  - RoBERTa:    peyterho/financial-roberta-large-macro-sentiment
+  - ClimateBERT: peyterho/climatebert-macro-sentiment
 """
 
 import re
@@ -20,6 +25,18 @@ from transformers import (
 )
 
 
+# ─── Default model names ─────────────────────────────────────────
+# Fine-tuned on 20K combined financial sentiment corpus (v0.2.0)
+DEFAULT_FINBERT = "peyterho/finbert-macro-sentiment"
+DEFAULT_ROBERTA = "peyterho/financial-roberta-large-macro-sentiment"
+DEFAULT_CLIMATEBERT = "peyterho/climatebert-macro-sentiment"
+
+# Original off-the-shelf models (v0.1.0)
+ORIGINAL_FINBERT = "ProsusAI/finbert"
+ORIGINAL_ROBERTA = "soleimanian/financial-roberta-large-sentiment"
+ORIGINAL_CLIMATEBERT = "climatebert/distilroberta-base-climate-sentiment"
+
+
 class SentimentHead:
     """Base class for a transformer sentiment scoring head."""
     
@@ -67,8 +84,14 @@ class SentimentHead:
 
 
 class FinBERTHead(SentimentHead):
-    """FinBERT for financial news sentiment. Accepts custom model_name for fine-tuned variants."""
-    def __init__(self, model_name="ProsusAI/finbert", device="cpu"):
+    """FinBERT for financial news sentiment.
+    
+    Args:
+        model_name: HF model ID. Default: fine-tuned 'peyterho/finbert-macro-sentiment'.
+                    Use 'ProsusAI/finbert' for the original off-the-shelf model.
+        device: 'cpu' or 'cuda:0'.
+    """
+    def __init__(self, model_name=DEFAULT_FINBERT, device="cpu"):
         super().__init__(
             model_name=model_name,
             label_map={0: "positive", 1: "negative", 2: "neutral"},
@@ -78,19 +101,33 @@ class FinBERTHead(SentimentHead):
 
 
 class FinancialRoBERTaHead(SentimentHead):
-    def __init__(self, device="cpu"):
+    """Financial-RoBERTa-Large for policy/formal text sentiment.
+    
+    Args:
+        model_name: HF model ID. Default: fine-tuned 'peyterho/financial-roberta-large-macro-sentiment'.
+                    Use 'soleimanian/financial-roberta-large-sentiment' for the original.
+        device: 'cpu' or 'cuda:0'.
+    """
+    def __init__(self, model_name=DEFAULT_ROBERTA, device="cpu"):
         super().__init__(
-            model_name="soleimanian/financial-roberta-large-sentiment",
-            label_map={0: "positive", 1: "negative", 2: "neutral"},
+            model_name=model_name,
+            label_map={0: "negative", 1: "neutral", 2: "positive"},
             score_map={"positive": 1.0, "negative": -1.0, "neutral": 0.0},
             device=device, max_length=512,
         )
 
 
 class ClimateBERTHead(SentimentHead):
-    def __init__(self, device="cpu"):
+    """ClimateBERT for climate risk/opportunity sentiment.
+    
+    Args:
+        model_name: HF model ID. Default: fine-tuned 'peyterho/climatebert-macro-sentiment'.
+                    Use 'climatebert/distilroberta-base-climate-sentiment' for the original.
+        device: 'cpu' or 'cuda:0'.
+    """
+    def __init__(self, model_name=DEFAULT_CLIMATEBERT, device="cpu"):
         super().__init__(
-            model_name="climatebert/distilroberta-base-climate-sentiment",
+            model_name=model_name,
             label_map={0: "risk", 1: "neutral", 2: "opportunity"},
             score_map={"risk": -1.0, "neutral": 0.0, "opportunity": 1.0},
             device=device, max_length=512,
@@ -169,13 +206,50 @@ class TopicRouter:
 
 
 class TransformerEnsemble:
-    def __init__(self, device="cpu", use_router=True, finbert_model="ProsusAI/finbert"):
+    """Multi-head transformer ensemble with domain routing.
+    
+    Args:
+        device: 'cpu' or 'cuda:0'.
+        use_router: Enable keyword-based topic routing.
+        finbert_model: HF model ID for the FinBERT head.
+        roberta_model: HF model ID for the Financial-RoBERTa head.
+        climatebert_model: HF model ID for the ClimateBERT head.
+    
+    Default models are fine-tuned on 20K combined financial corpus.
+    Pass the ORIGINAL_* constants to use off-the-shelf models.
+    
+    Examples:
+        # Use fine-tuned models (default, recommended)
+        ensemble = TransformerEnsemble(device="cpu")
+        
+        # Use original off-the-shelf models
+        from macro_sentiment.transformers_ensemble import ORIGINAL_FINBERT, ORIGINAL_ROBERTA, ORIGINAL_CLIMATEBERT
+        ensemble = TransformerEnsemble(
+            finbert_model=ORIGINAL_FINBERT,
+            roberta_model=ORIGINAL_ROBERTA,
+            climatebert_model=ORIGINAL_CLIMATEBERT,
+        )
+        
+        # Mix and match
+        ensemble = TransformerEnsemble(
+            finbert_model="peyterho/finbert-macro-sentiment",
+            roberta_model="soleimanian/financial-roberta-large-sentiment",
+        )
+    """
+    def __init__(
+        self,
+        device="cpu",
+        use_router=True,
+        finbert_model=DEFAULT_FINBERT,
+        roberta_model=DEFAULT_ROBERTA,
+        climatebert_model=DEFAULT_CLIMATEBERT,
+    ):
         self.device = device
         self.use_router = use_router
         self.heads = {
             "finbert": FinBERTHead(model_name=finbert_model, device=device),
-            "policy": FinancialRoBERTaHead(device),
-            "climate": ClimateBERTHead(device),
+            "policy": FinancialRoBERTaHead(model_name=roberta_model, device=device),
+            "climate": ClimateBERTHead(model_name=climatebert_model, device=device),
         }
         self.router = TopicRouter(device) if use_router else None