| | 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) |
| |
|
