|
|
from sklearn.metrics import accuracy_score, f1_score |
|
|
from sklearn.linear_model import LogisticRegression |
|
|
import datasets |
|
|
import numpy as np |
|
|
|
|
|
import torch |
|
|
from llm2vec import LLM2Vec |
|
|
|
|
|
dataset = "mteb/amazon_counterfactual" |
|
|
instruction = "Classify a given Amazon customer review text as either counterfactual or notcounterfactual: " |
|
|
|
|
|
dataset = datasets.load_dataset(dataset, "en") |
|
|
|
|
|
sentences_train, y_train = dataset["train"]["text"], dataset["train"]["label"] |
|
|
sentences_test, y_test = dataset["test"]["text"], dataset["test"]["label"] |
|
|
max_iter = 100 |
|
|
batch_size = 8 |
|
|
|
|
|
scores = {} |
|
|
clf = LogisticRegression( |
|
|
random_state=42, |
|
|
n_jobs=1, |
|
|
max_iter=max_iter, |
|
|
verbose=0, |
|
|
) |
|
|
|
|
|
print("Loading model...") |
|
|
model = LLM2Vec.from_pretrained( |
|
|
"McGill-NLP/LLM2Vec-Mistral-7B-Instruct-v2-mntp", |
|
|
peft_model_name_or_path="McGill-NLP/LLM2Vec-Mistral-7B-Instruct-v2-mntp-supervised", |
|
|
device_map="cuda" if torch.cuda.is_available() else "cpu", |
|
|
torch_dtype=torch.bfloat16, |
|
|
) |
|
|
|
|
|
|
|
|
def append_instruction(instruction, sentences): |
|
|
new_sentences = [] |
|
|
for s in sentences: |
|
|
new_sentences.append([instruction, s, 0]) |
|
|
return new_sentences |
|
|
|
|
|
|
|
|
print(f"Encoding {len(sentences_train)} training sentences...") |
|
|
sentences_train = append_instruction(instruction, sentences_train) |
|
|
X_train = np.asarray(model.encode(sentences_train, batch_size=batch_size)) |
|
|
|
|
|
print(f"Encoding {len(sentences_test)} test sentences...") |
|
|
sentences_test = append_instruction(instruction, sentences_test) |
|
|
X_test = np.asarray(model.encode(sentences_test, batch_size=batch_size)) |
|
|
|
|
|
print("Fitting logistic regression classifier...") |
|
|
clf.fit(X_train, y_train) |
|
|
print("Evaluating...") |
|
|
y_pred = clf.predict(X_test) |
|
|
|
|
|
accuracy = accuracy_score(y_test, y_pred) |
|
|
scores["accuracy"] = accuracy |
|
|
f1 = f1_score(y_test, y_pred, average="macro") |
|
|
scores["f1"] = f1 |
|
|
|
|
|
print(scores) |
|
|
|
|
|
|
