---
tags:
- pytorch_model_hub_mixin
- model_hub_mixin
---

This model adds a classification head on top of [LofiAmazon/BarcodeBERT-Entire-BOLD](https://huggingface.co/LofiAmazon/BarcodeBERT-Entire-BOLD), The classification head is a linear layer which concatenates the DNA embeddings with environmental layer data. This model has been trained with 
[BOLD-Embeddings-Ecolayers-Amazon](https://huggingface.co/datasets/LofiAmazon/BOLD-Embeddings-Ecolayers-Amazon) to predict taxonomic genuses. BOLD-Embeddings-Ecolayers-Amazon includes DNA embeddings and ecological raster layer data regarding sample location.

## Example Usage

First you need to download the `StandardScaler` used to normalise environmental layer values during training. You can find it [here](https://huggingface.co/spaces/LofiAmazon/LofiAmazonSpace/blob/main/scaler.pkl).

You will also need to download all ecological layers from [here](https://huggingface.co/datasets/LofiAmazon/Global-Ecolayers/tree/main).

The model will output a probability distribution over genera that were present in the training dataset. You can find the mapping between the index in the vector and the genus name in [this file](https://huggingface.co/spaces/LofiAmazon/LofiAmazonSpace/blob/main/genus_labels.json).

```py
import pickle

from transformers import PreTrainedTokenizerFast, BertForMaskedLM, BertConfig
import rasterio
from rasterio.sample import sample_gen


class DNASeqClassifier(nn.Module, PyTorchModelHubMixin):
    def __init__(self, bert_model, env_dim, num_classes):
        super(DNASeqClassifier, self).__init__()
        self.bert = bert_model
        self.env_dim = env_dim
        self.num_classes = num_classes
        self.fc = nn.Linear(768 + env_dim, num_classes)

    def forward(self, bert_inputs, env_data):
        outputs = self.bert(**bert_inputs)
        dna_embeddings = outputs.hidden_states[-1].mean(1)
        combined = torch.cat((dna_embeddings, env_data), dim=1)
        logits = self.fc(combined)

        return logits

classification_model = DNASeqClassifier.from_pretrained(
    "LofiAmazon/BarcodeBERT-Finetuned-Amazon",
    bert_model=BertForMaskedLM(
        BertConfig(vocab_size=259, output_hidden_states=True),
    ),
)

ecolayers = [
    "median_elevation_1km.tiff",
    "human_footprint.tiff",
    "population_density_1km.tif",
    "annual_precipitation.tif",
    "precipitation_seasonality.tif",
    "annual_mean_air_temp.tif",
    "temp_seasonality.tif",
]

with open("scaler.pkl", "rb") as f:
    scaler = pickle.load(f)

tokenizer = PreTrainedTokenizerFast.from_pretrained("LofiAmazon/BarcodeBERT-Entire-BOLD")

# The DNA sequence you want to predict.
# There should be a space after every 4 characters.
# The sequence may also have unknown characters which are not A,C,T,G.
# The maximum DNA sequence length (not counting spaces) should be 660 characters
dna_sequence = "AACA ATGT ATTT A-T- TTCG CCCT TGTG AATT TATT ..."
# Location where DNA was sampled
coords = (-3.009083, -58.68281)

# Tokenize the DNA sequence
dna_tokenized = tokenizer(dna_sequence, return_tensors="pt")

# Obtain the environmental data from the coordinates
env_data = []
for layer in ecolayers:
    with rasterio.open(layer) as dataset:
        # Get the corresponding ecological values for the samples
        results = sample_gen(dataset, [coords])
        results = [r for r in results]
    layer_data = np.mean(results[0])
    env_data.append(layer_data)
env_data = scaler.transform([env_data])
env_data = torch.from_numpy(env_data).to(torch.float32)

# Obtain genus prediction
logits = classification_model(dna_tokenized, env_data)
temperature = 0.2

# Obtain the final genus probabilities
probs = torch.softmax(logits / temperature, dim=1).squeeze()
```