File size: 1,899 Bytes
e9cd0c7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
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)
# {'accuracy': 0.891044776119403, 'f1': 0.8283106625713033}
|