Tiny Transformers for Aspect-Based Financial Sentiment Classification

"""
# TASK DESCRIPTION

Read through the following snippet of text carefully and classify the **financial sentiment** as
either negative, neutral, or positive. You must also provide a short rationale for why you
assigned the financial sentiment you did.

For clarity here are the definitions negative, neutral, and positive sentiments:

   - **Negative**: The snippet describes an event or development that has had, is having, or is
       expected to have a material negative impact on the company’s financial performance, share price, reputation,
       or outlook.

   - **Neutral**: The snippet is informational/descriptive and is not expected to have a material positive or
       negative impact on the company.

   - **Positive**: The snippet describes an event or development that has had, is having, or is expected to have
       a material positive impact on the company’s financial performance, share price, reputation, or outlook.

Materiality note:
- “Material impact” includes likely effects on share price, revenue, costs, profitability, cash flow, guidance,
regulatory exposure, reputation, risk exposure, or competitive position.

# TASK GUIDELINES

For the avoidance of doubt here is a decision tree that you can follow to arrive at the most appropriate
sentiment classification for the snippet. Pay careful attention to the logic. Don't deviate.

START
│
├── Step 1: Carefully read and understand the snippet.
│
├── Step 2: Check for sentiment indicators:
│
├── Is the snippet clearly NEGATIVE?
│     (share price decline, losses, scandals, lawsuits, layoffs, product recalls, regulatory fines,
│      leadership resignations, declining sales, market-share losses, reputational damage etc.)
│    │
│    ├── YES → Classify as "negative"
│    │      └── Provide a rationale by summarizing WHY the snippet is negative.
│    │
│    └── NO → Continue below
│
├── Is the snippet clearly POSITIVE?
│     (share price increases, strong earnings, favorable partnerships, successful product launches, awards,
│      expansion plans, positive analyst coverage, reputational enhancement, etc.)
│    │
│    ├── YES → Classify as "positive"
│    │      └── Provide a rationale by summarizing WHY the snippet is positive.
│    │
│    └── NO → Continue below
│
└── If neither clearly positive nor negative → Classify as "neutral"
         (routine product announcements without performance implications, leadership appointments,
          scheduled reports, factual statements, general industry overviews, etc.)
       └── Provide a rationale by summarizing the WHY the snippet is neutral.

If there is conflicting sentiment in the snippet pick the most dominant one, otherwise default to **neutral**.

# RESPONSE FORMAT

You must respond with ONLY a valid JSON object formatted as follows. DO NOT WRITE ANY PREAMBLE JUST RETURN JSON.

{{
   "rationale": "A one-sentence rationale for your classification",
   "financial_sentiment": "either negative, neutral, or positive"
}}

# SNIPPET TO LABEL

Here is the snippet we would like you to assign a negative, neutral, or positive financial sentiment label to:

{text}

P.S. REMEMBER TO READ THE SNIPPET CAREFULLY AND FOLLOW THE GUIDELINES TO ARRIVE AT THE MOST APPROPRIATE
FINANCIAL SENTIMENT CLASSIFICATION. WHEN IN DOUBT, YOU SHOULD DEFER TO A "neutral" CLASSIFICATION FOR THE SNIPPET.
GOOD LUCK!
"""

import os
import datetime as dt
import polars as pl
from sklearn.linear_model import SGDClassifier
from sklearn.metrics import confusion_matrix, accuracy_score

def embedding_model(data: pl.DataFrame, model: str, embeddings_path: str) -> tuple:
    """
    Train a linear classifier on text embeddings for financial sentiment classification.
    
    Generates or loads embeddings, trains an SGDClassifier with hinge loss, and evaluates
    performance on train/validation splits. Returns predictions, accuracy metrics, and
    confusion matrices for both splits.
    
    :param data: DataFrame containing 'text' and 'majority_vote' columns with samples and labels
    :param model: Embedding model identifier for OpenRouter API
    :param embeddings_path: Path to cache/load embeddings (created if doesn't exist)
    :return: Returns a tuple of predictions, accuracy metrics, and confusion matrices for both splits.
    """
    # Memoize the embeddings for subsequent training runs.
    if not os.path.exists(embeddings_path):
        # Generate the embeddings using OpenRouter.
        embed_data(data=data, out_file=embeddings_path, model=model)

    # Load the embeddings, only as many samples as this training run.
    embeddings = np.load(embeddings_path)
    embeddings = embeddings[:data.shape[0]]

    # Fetch labels using majority vote and map to model targets.
    label_to_target={"positive": 1, "neutral": 0, "negative": -1}
    labels = list(label_to_target.keys())
    targets = [label_to_target[row['majority_vote']] for row in data.to_dicts()]

    # Samples already shuffled, simple subscript split is sufficient.
    train_pct = 0.8
    n_train = int(data.shape[0] * train_pct)
    x_train, y_train = embeddings[:n_train], targets[:n_train]
    x_val, y_val = embeddings[n_train:], targets[n_train:]

    # Get the training and validation text.
    train_text = data[:n_train].select("text").to_numpy().flatten().tolist()
    val_text = data[n_train:].select("text").to_numpy().flatten().tolist()

    classifier = SGDClassifier(
        loss='hinge',
        penalty='l2',
        alpha=0.00001,
        validation_fraction=0.00001
    )
    classifier.fit(x_train, y_train)

    # Make predictions.
    train_preds = classifier.predict(x_train)
    val_preds = classifier.predict(x_val)

    # Score the preds.
    train_acc = accuracy_score(y_true=y_train, y_pred=train_preds)
    val_acc = accuracy_score(y_true=y_val, y_pred=val_preds)

    confusion_train = confusion_matrix(y_true=y_train, y_pred=train_preds)
    confusion_train = {
        r: {
            c: int(confusion_train[i, j])
            for j, c in enumerate(labels)
        }
        for i, r in enumerate(labels)
    }
    confusion_val = confusion_matrix(y_true=y_val, y_pred=val_preds)
    confusion_val = {
        r: {
            c: int(confusion_val[i, j])
            for j, c in enumerate(labels)
        }
        for i, r in enumerate(labels)
    }

    print(f"{dt.datetime.utcnow()} - Train Accuracy {train_acc:.4f}")
    print(f"{dt.datetime.utcnow()} - Val Accuracy {val_acc:.4f}")
    print(f"{dt.datetime.utcnow()} - Confusion Matrix")
    print("Train", simplejson.dumps(confusion_train, indent=4))
    print("Val", simplejson.dumps(confusion_val, indent=4))

    return (
        (y_train, train_preds, train_text),
        (y_val, val_preds, val_text),
        (train_acc, val_acc),
        (confusion_train, confusion_val)
    )

"""
Consult the oracle to evaluate whether it agrees with the linear models prediction, or if we need to relabel this text sample.
"""
import textwrap

labels = ["negative", "neutral", "positive"]
labelling_prompt = textwrap.dedent(financial_sentiment_prompt(text))
label_options = " | ".join([f'"{label}"' for label in labels])
# The prediction variable is one of the labels above, where all linear models agreed.

messages = [
    {
        "role": "system",
        "content": [
            {
                "type": "text",
                "text":  textwrap.dedent(
                    f"""
                    # TASK DESCRIPTION
                    
                    Given the following LLM prompt, another model labelled this text as \"{prediction}\".
                    Is this model correct? Provide the correct label and a rationale for your answer.
                    
                    # RESPONSE FORMAT
                    
                    You must respond using this exact format. DO NOT RESPOND IN ANY OTHER WAY:

                    ```json
                    {{
                        "correct": "true" | "false",
                        "label": {label_options},
                        "reason": "A rationale for the correctness or incorrectness of your label."
                    }}
                    ```
                    
                    P.S. Think fast there is no time to waste.                        
                    """
                )
            }
        ]
    },
    {
        "role": "user",
        "content": [{ "type": "text", "text": labelling_prompt }]
    }
]

# Send these messages to OpenRouter using the OpenAI client.

import os
import datetime as dt
import random
import textwrap
from itertools import chain

import polars as pl
import torch
from datasets import Dataset, DatasetDict
from sklearn.metrics import f1_score, accuracy_score
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    DataCollatorForSeq2Seq,
    TrainingArguments,
    Trainer, AutoModelForSequenceClassification, DataCollatorWithPadding,
)
from functools import partial


def eval_finbert():
    """
    Evaluate the finbert model.

    :return:
    """
    import numpy as np

    def compute_metrics(eval_pred):
        logits, labels = eval_pred
        predictions = np.argmax(logits, axis=-1)

        # Prefer Macro F1.
        f1 = f1_score(labels, predictions, average="macro")
        accuracy = accuracy_score(labels, predictions)

        return {"f1": f1, "accuracy": accuracy}

    finbert_model_id = "ProsusAI/finbert"

    finbert_label2id = {
        "positive": 0,
        "negative": 1,
        "neutral": 2,
    }
    finbert_id2label = {v: k for k, v in finbert_label2id.items()}
    finbert_num_labels = len(finbert_label2id)

    finbert_tokenizer = AutoTokenizer.from_pretrained(finbert_model_id)
    finbert_model = AutoModelForSequenceClassification.from_pretrained(
        finbert_model_id,
        num_labels=finbert_num_labels,
        label2id=finbert_label2id,
        id2label=finbert_id2label,
    )

    def tokenize_finbert(batch):
        return finbert_tokenizer(batch["text"], truncation=True, max_length=512)

    # PhraseBank with FinBERT tokenizer.
    finbert_phrase_ds = fin_bank_ds.map(tokenize_finbert, batched=True)
    val_ds_for_finbert = val_ds.map(tokenize_finbert, batched=True)

    finbert_collator = DataCollatorWithPadding(tokenizer=finbert_tokenizer)

    finbert_eval_args = TrainingArguments(
        output_dir="finbert_baseline_eval",
        per_device_eval_batch_size=32,
        do_train=False,
        do_eval=True,
        logging_strategy="no",
    )

    # Evaluate FinBERT on PhraseBank.
    finbert_trainer_phrase = Trainer(
        model=finbert_model,
        args=finbert_eval_args,
        data_collator=finbert_collator,
        eval_dataset=finbert_phrase_ds,
        compute_metrics=compute_metrics,
    )

    print("\n---")
    print("Evaluating FinBERT on Financial PhraseBank...")
    finbert_phrase_metrics = finbert_trainer_phrase.evaluate()
    print("FinBERT metrics on PhraseBank:", finbert_phrase_metrics)
    print("---")

    # Evaluate FinBERT on 20% of data_labels.
    finbert_trainer_val = Trainer(
        model=finbert_model,
        args=finbert_eval_args,
        data_collator=finbert_collator,
        eval_dataset=val_ds_for_finbert,
        compute_metrics=compute_metrics,
    )

    print("\n---")
    print("Evaluating FinBERT on 20% held-out data_labels...")
    finbert_val_metrics = finbert_trainer_val.evaluate()
    print("FinBERT metrics on 20% data_labels held-out set:", finbert_val_metrics)
    print("---\n")

    # Move the finbert model off the gpu.
    finbert_model.to(device="cpu")


def compute_metrics_without_eos(eval_pred, eos_token_id):
    """
    Compute the f1 score and the accuracy based on TOKEN classification. Given we are predicting a single token
    this is a good approximation of the actual scores.

    Make sure to remove eos token.

    :param eval_pred: Predictions.
    :param eos_token_id: The eos token id.
    :return:
    """
    predictions, labels = eval_pred
    predictions = predictions[:, :-1]
    labels = labels[:, 1:]

    # Exclude Prompt (-100) AND EOS.
    valid_mask = (labels != -100) & (labels != eos_token_id)

    pred_flat = predictions[valid_mask]
    label_flat = labels[valid_mask]

    f1 = f1_score(label_flat, pred_flat, average="macro")
    accuracy = accuracy_score(label_flat, pred_flat)
    return {"f1": f1, "accuracy": accuracy}


def preprocess_logits_for_metrics(logits, labels):
    """
    Original logits are (Batch, Seq, Vocab).
    We only need (Batch, Seq) containing the indices of the max logit.
    """
    if isinstance(logits, tuple):
        # Depending on the model and config, logits may contain extra tensors,
        # like past_key_values, but logits always come first
        logits = logits[0]
    return logits.argmax(dim=-1)


if __name__ == "__main__":

    # ----------------------------------------------------
    # STEP 1 - Load financial sentiment dataset
    # ----------------------------------------------------

    task = "financial_sentiment"

    # Target TEXT strings for the Guard model.
    label_map = {
        "positive": "positive",
        "negative": "negative",
        "neutral": "neutral"
    }

    modeling_dir = os.path.join("/", "workspace", ".nosible")

    data_labels = pl.read_ipc(os.path.join(modeling_dir, f"{task}_100.0k_iter_18.ipc"))
    fin_bank = pl.read_ndjson(os.path.join(modeling_dir, "financial_phrase_bank.ndjson"))

    # Select text and text-labels.
    # Check if 'majority_vote' exists (it does in your IPC), otherwise rename or use 'labels'.
    if "majority_vote" in data_labels.columns:
        data = data_labels.select(["text", "majority_vote"]).rename({"majority_vote": "labels"})
    else:
        data = data_labels.select(["text", "labels"])

    fin_bank = fin_bank.select(["text", "labels"])

    print(f"Data shape: {data.shape}")

    # ----------------------------------------------------
    # STEP 2 - Create train/val splits
    # ----------------------------------------------------

    # Slice to 100k only if we have that much data, otherwise take all.
    limit = min(100_000, len(data))
    data = data[:limit]

    full_ds = Dataset.from_polars(df=data)
    fin_bank_ds = Dataset.from_polars(df=fin_bank)

    # Now this will succeed because we have at least 20 rows (16 train, 4 test)
    split_ds = full_ds.train_test_split(test_size=0.2, seed=42)
    train_ds = split_ds["train"]
    val_ds = split_ds["test"]

    ds = DatasetDict({
        "train": train_ds,
        "val": val_ds,
        "phrasebank": fin_bank_ds,
    })

    # ----------------------------------------------------
    # STEP 3 - Fine-tune Qwen 3 (Guard Approach)
    # ----------------------------------------------------

    # UPDATED: Use Qwen 3 (0.6B).
    model_id = "Qwen/Qwen3-0.6B"

    # Important: trust_remote_code=True is essential for Qwen3.
    # Make sure to set left padding.
    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 = "left"

    # Load Model.
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        trust_remote_code=True,
        dtype=torch.bfloat16,
        device_map="auto"
    )

    def tokenize(batch):
        """
        Tokenize a batch.

        :param batch: Batch to tokenize.
        :return:
        """

        input_ids_list = []
        attention_mask_list = []
        labels_list = []

        for text, label in zip(batch['text'], batch['labels']):

            # 1. Build the prompt (User part)
            system = "Classify the financial sentiment as positive, negative, or neutral."
            msgs = [
                {"role": "system", "content": system},
                {"role": "user", "content": text},
            ]
            prompt_str = tokenizer.apply_chat_template(
                msgs,
                tokenize=False,
                add_generation_prompt=True,
                enable_thinking=False, # By default, this is true.
            )

            # 2. Build the answer (Assistant part)
            answer_str = f"{label}<|im_end|>" # Standard Qwen end token.

            # 3. Tokenize separately to know lengths
            prompt_tokens = tokenizer(prompt_str, add_special_tokens=False)
            answer_tokens = tokenizer(answer_str, add_special_tokens=False)

            assert len(answer_tokens["input_ids"]) == 2

            # 4. Combine
            input_ids = prompt_tokens["input_ids"] + answer_tokens["input_ids"]
            attention_mask = prompt_tokens["attention_mask"] + answer_tokens["attention_mask"]

            # 5. CREATE LABELS WITH MASKING
            # -100 tells PyTorch to IGNORE these tokens during training.
            labels = [-100] * len(prompt_tokens["input_ids"]) + answer_tokens["input_ids"]

            # Truncate if necessary (simplified for brevity).
            if len(input_ids) > 2048:
                input_ids = input_ids[:2048]
                attention_mask = attention_mask[:2048]
                labels = labels[:2048]

            input_ids_list.append(input_ids)
            attention_mask_list.append(attention_mask)
            labels_list.append(labels)

        return {
            "input_ids": input_ids_list,
            "attention_mask": attention_mask_list,
            "labels": labels_list
        }

    # Tokenize the datasets.
    t_ds = ds.map(tokenize, batched=True, remove_columns=ds["train"].column_names, num_proc=8)

    # Set the collator for padding.
    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, padding=True)

    timestamp = dt.datetime.now().strftime('%Y%m%d_%H%M%S')
    task_output = f"{task}_qwen3_0.6B_{timestamp}"

    training_args = TrainingArguments(
        output_dir=task_output,

        # Qwen3 0.6B is very small, we might be able to increase batch size slightly.
        gradient_accumulation_steps=4,
        per_device_train_batch_size=16,
        per_device_eval_batch_size=16,

        num_train_epochs=5,
        learning_rate=2e-5,
        lr_scheduler_type="cosine",
        warmup_ratio=0.03,
        max_grad_norm=1.0,
        weight_decay=0.1,

        # Improvements for instruction finetuning.
        neftune_noise_alpha=5,

        # Optimizations.
        bf16=True,
        optim="adamw_torch_fused",
        logging_strategy="steps",
        logging_steps=10,
        logging_dir=task_output,
        logging_first_step=True,
        group_by_length=True,
        torch_compile=False,

        # Eval params.
        eval_strategy="steps",
        eval_steps=500,
        save_total_limit=3,
        load_best_model_at_end=True,
        metric_for_best_model="eval_val_loss",
        report_to="none",
        save_strategy="steps",
        save_steps=500,

    )

    n_eval_train = int(0.05 * len(t_ds["train"]))
    chat_end_token_id = tokenizer.convert_tokens_to_ids("<|im_end|>")

    trainer = Trainer(
        model=model,
        args=training_args,
        data_collator=data_collator,
        train_dataset=t_ds["train"],
        eval_dataset={
            "train": t_ds["train"].select(list(range(n_eval_train))),
            "val": t_ds["val"],
            "phrasebank": t_ds["phrasebank"],
        },
        compute_metrics=partial(compute_metrics_without_eos, eos_token_id=chat_end_token_id),
        preprocess_logits_for_metrics=preprocess_logits_for_metrics,
    )

    trainer.train()

import math
from openai import OpenAI

# Initialize the client pointing to your vLLM server
client = OpenAI(
    base_url="YOUR_ENDPOINT_URL_HERE/v1",
    api_key="YOUR_API_KEY_HERE"
)

model_id = "NOSIBLE/financial-sentiment-v1.1-base"

# Input text to classify
text = "The company reported a record profit margin of 15% this quarter."

# Define the classification labels
labels = ["positive", "negative", "neutral"]

# Prepare the conversation
messages = [
    {"role": "system", "content": "Classify the financial sentiment as positive, neutral, or negative."},
    {"role": "user", "content": text},
]

# Make the API call
chat_completion = client.chat.completions.create(
    model=model_id,
    messages=messages,
    temperature=0,
    max_tokens=1,
    stream=False,
    logprobs=True,              # Enable log probabilities to calculate confidence
    top_logprobs=len(labels),   # Ensure we capture logprobs for our choices
    extra_body={
        "chat_template_kwargs": {"enable_thinking": False}, # Must be set to false.
        "regex": "(positive|neutral|negative)",
    },
)

# Extract the response content
response_label = chat_completion.choices[0].message.content

# Extract the logprobs for the generated token to calculate confidence
first_token_logprobs = chat_completion.choices[0].logprobs.content[0].top_logprobs

print(f"--- Classification Results ---")
print(f"Input: {text}")
print(f"Predicted Label: {response_label}\n")

print("--- Label Confidence ---")
for lp in first_token_logprobs:
    # Convert log probability to percentage
    probability = math.exp(lp.logprob)
    print(f"Token: '{lp.token}' | Probability: {probability:.2%}")