CutiesNLP

What is the purpose of the final layers?

Which layers provide the most useful representations for this task?

Extract hidden states from different layers and use them as features for a logistic regression classifier.

import numpy as np
import torch
from sklearn.model_selection import StratifiedKFold
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline


def get_hs(model, train, test):
    Xs = []
    y = []
    y2 = []

    for idx in range(len(train)):
        data = train[idx]
        pos = data["eol_positions"]
        cr = data["correct_idx"]
        embs = torch.from_numpy(data["embeds_after_l0"]).unsqueeze(0).cuda().half()
        with torch.no_grad():
            out = model.model(inputs_embeds=embs, output_hidden_states=True)
        hs = torch.cat(out.hidden_states)
        for i in range(4):
            Xs.append(hs[:, pos[i]].cpu().numpy())
            y.append(int(i == cr))
        y2.append(cr)

    Xs = np.stack(Xs).reshape(len(Xs) // 4, 4, 24, -1)
    y = np.array(y).reshape(len(y) // 4, 4)
    y2 = np.array(y2)

    Xs_test = []

    for idx in range(len(test)):
        data = test[idx]
        pos = data["eol_positions"]
        embs = torch.from_numpy(data["embeds_after_l0"]).unsqueeze(0).cuda().half()
        with torch.no_grad():
            out = model.model(inputs_embeds=embs, output_hidden_states=True)
        hs = torch.cat(out.hidden_states)
        for i in range(4):
            Xs_test.append(hs[:, pos[i]].cpu().numpy())

    Xs_test = np.stack(Xs_test).reshape(len(Xs_test) // 4, 4, 24, -1)

    return Xs, y, y2, Xs_test


def evaluate(Xs, y, y2, ly):
    X_layer = Xs[:, :, ly]

    train_accs = []
    val_accs = []

    skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=0)

    for train_idx, val_idx in skf.split(X_layer, y2):
        X_train = X_layer[train_idx].reshape(4 * len(train_idx), -1)
        X_val = X_layer[val_idx].reshape(4 * len(val_idx), -1)

        y_train = y[train_idx].reshape(-1)

        clf = Pipeline([
            ("scaler", StandardScaler()),
            ("logreg", LogisticRegression(random_state=0, max_iter=1000, C=1e-2)),
        ])

        clf.fit(X_train, y_train)

        pred = clf.predict_proba(X_train)[:, 1]
        train_acc = (pred.reshape(-1, 4).argmax(axis=1) == y2[train_idx]).mean()

        pred = clf.predict_proba(X_val)[:, 1]
        val_acc = (pred.reshape(-1, 4).argmax(axis=1) == y2[val_idx]).mean()

        train_accs.append(train_acc)
        val_accs.append(val_acc)

    return np.mean(train_accs), np.mean(val_accs)


def experiment(Xs, y, y2):
    for ly in range(24):
        train_acc, val_acc = evaluate(Xs, y, y2, ly)

        print(f"Layer {ly}" f" {train_acc:.3f}" f" {val_acc:.3f}")

    for ly in [[13, 14], [13, 15], [14, 15], [13, 14, 15]]:
        train_acc, val_acc = evaluate(Xs, y, y2, ly)

        print(f"Layer {ly}" f" {train_acc:.3f}" f" {val_acc:.3f}")


def solution(model, train, test):
    Xs, y, y2, Xs_test = get_hs(model, train, test)

    ly = [13, 14]

    X = Xs[:, :, ly]
    X = X.reshape(len(X) * 4, -1)

    clf = Pipeline([
        ("scaler", StandardScaler()),
        ("logreg", LogisticRegression(random_state=0, max_iter=1000, C=1e-2)),
    ]).fit(X, y.reshape(-1))

    X_test = Xs_test[:, :, ly]
    X_test = X_test.reshape(len(X_test) * 4, -1)

    test_preds = clf.predict_proba(X_test)[:, 1].reshape(-1, 4).argmax(axis=1)

    return test_preds

Download solution.py