Knots ProtBert-BFD AlphaFold

Fine-tuned ProtBert-BFD to classify proteins as knotted vs. unknotted.

Model Details

• Model type: Bert
• Language: proteins (amino acid sequences)
• Finetuned from model: Rostlab/prot_bert_bfd

Model Sources:

• Repository: CEITEC
• Paper: TBD

Usage

Dataset format:

id,sequence,label
A0A2W5F4Z7,MGGIFRVNTYYTDLEPYLQSTKLPIYGALLDGENIYELVDKSKGILVIGNESKGIRSTIQNFIQKPITIPRIGQAESLNAAVATGIIVGQLTL,1
...

Load the dataset:

import pandas as pd
from datasets import Dataset, load_dataset

df = pd.read_csv(INPUT, sep=',')
dss = Dataset.from_pandas(df)

Predict:

import torch
import numpy as np
from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments
from math import exp

def tokenize_function(s):
    seq_split = ' '.join(s['Sequence'])
    return tokenizerM1(seq_split)

tokenizer = AutoTokenizer.from_pretrained('roa7n/knots_protbertBFD_alphafold')
model = AutoModelForSequenceClassification.from_pretrained('roa7n/knots_protbertBFD_alphafold')

tokenized_dataset = dss.map(tokenize_function, num_proc=4)
tokenized_dataset.set_format('pt')
tokenized_dataset

training_args = TrainingArguments(<PATH>, fp16=True, per_device_eval_batch_size=50, report_to='none')  

trainer = Trainer(
    model,
    training_args,
    train_dataset=tokenized_dataset,
    eval_dataset=tokenized_dataset,
    tokenizer=tokenizerM1
)

predictions, _, _ = trainer.predict(tokenized_dataset)
predictions = [np.exp(p[1]) / np.sum(np.exp(p), axis=0) for p in predictions]
df['preds'] = predictions

Evaluation

Per protein family metrics:

Citation [optional]

BibTeX: TODO

Model Authors

Simecek: simecek@mail.muni.cz Klimentova: vae@mail.muni.cz Sramkova: denisa.sramkova@mail.muni.cz