|
import argparse |
|
import torch |
|
from transformers import AutoModelForSequenceClassification, AutoTokenizer |
|
from datasets import Dataset |
|
import pandas as pd |
|
import numpy as np |
|
from sklearn.metrics import accuracy_score, f1_score, matthews_corrcoef, balanced_accuracy_score, roc_auc_score, confusion_matrix |
|
from utils import fasta_to_df |
|
|
|
def compute_metrics(logits, labels): |
|
predictions = np.argmax(logits, axis=1) |
|
labels = np.array(labels, dtype=int) |
|
predictions = np.array(predictions, dtype=int) |
|
|
|
acc = accuracy_score(labels, predictions) |
|
f1 = f1_score(labels, predictions, average='weighted') |
|
mcc = matthews_corrcoef(labels, predictions) |
|
balanced_acc = balanced_accuracy_score(labels, predictions) |
|
auc_roc = None |
|
|
|
if len(np.unique(labels)) == 2: |
|
probs = np.exp(logits[:, 1]) / np.sum(np.exp(logits), axis=1) |
|
auc_roc = roc_auc_score(labels, probs) |
|
|
|
cm = confusion_matrix(labels, predictions) |
|
return { |
|
'accuracy': acc, |
|
'f1_score': f1, |
|
'mcc': mcc, |
|
'auc_roc': auc_roc, |
|
'balanced_accuracy': balanced_acc, |
|
'confusion_matrix': cm.tolist() |
|
} |
|
|
|
def encode_sequence(sequence, tokenizer, max_length): |
|
return tokenizer(sequence, truncation=True, padding='max_length', max_length=max_length, return_tensors='pt') |
|
|
|
def evaluate(model_path, test_file=None, sequence=None): |
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
model = AutoModelForSequenceClassification.from_pretrained(model_path, ignore_mismatched_sizes=True).to(device) |
|
tokenizer = AutoTokenizer.from_pretrained(model_path) |
|
|
|
if sequence: |
|
inputs = encode_sequence(sequence, tokenizer, tokenizer.model_max_length) |
|
with torch.no_grad(): |
|
outputs = model(**{k: v.to(device) for k, v in inputs.items()}) |
|
logits = outputs.logits.cpu().numpy() |
|
print("Single Sequence Prediction:", np.argmax(logits, axis=1)) |
|
return |
|
|
|
test_df = fasta_to_df(test_file) |
|
label_map = { |
|
'non-pathogen': 0, |
|
'pathogen': 1 |
|
} |
|
test_df['label'] = test_df['label'].str.lower().map(label_map) |
|
dataset = Dataset.from_pandas(test_df) |
|
dataset = dataset.map(lambda x: encode_sequence(x['sequence'], tokenizer, tokenizer.model_max_length), batched=True) |
|
dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'label']) |
|
|
|
dataloader = torch.utils.data.DataLoader(dataset, batch_size=8) |
|
logits_list, labels_list = [], [] |
|
|
|
model.eval() |
|
with torch.no_grad(): |
|
for batch in dataloader: |
|
inputs = {k: v.to(device) for k, v in batch.items() if k != 'label'} |
|
|
|
labels = np.array(batch['label']) |
|
outputs = model(**inputs) |
|
logits_list.append(outputs.logits.cpu().numpy()) |
|
labels_list.append(labels) |
|
|
|
logits = np.concatenate(logits_list, axis=0) |
|
labels = np.concatenate(labels_list, axis=0) |
|
results = compute_metrics(logits, labels) |
|
print("Evaluation Metrics:", results) |
|
|
|
if __name__ == "__main__": |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument("--model_path", type=str, required=True, help="Path to the fine-tuned model directory") |
|
parser.add_argument("--test_file", type=str, help="Path to the test fasta file") |
|
parser.add_argument("--sequence", type=str, help="Single DNA sequence to classify") |
|
args = parser.parse_args() |
|
evaluate(args.model_path, args.test_file, args.sequence) |
|
|
