import torch
from transformers import AutoTokenizer, AutoModel, AutoConfig
from Bio import SeqIO
from sklearn.manifold import TSNE
import numpy as np
from tqdm import tqdm  # For progress tracking

# Device configuration
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

config = AutoConfig.from_pretrained('zhihan1996/DNABERT-S', trust_remote_code = True)

# Load model and tokenizer
tokenizer_16S = AutoTokenizer.from_pretrained('OceanOmics/eDNABERT-S_16S', trust_remote_code=True)
model_16S = AutoModel.from_pretrained('OceanOmics/eDNABERT-S_16S', trust_remote_code=True, config=config)
model_16S.to(device)
model_16S.eval()
names_12, seqs_12 = [], []
for seq in SeqIO.parse('16S_all_ASVs.fasta', 'fasta'):
        seqs_12.append(str(seq.seq))
        names_12.append(str(seq.id))

print(f"Processing {len(seqs_12)} sequences")


# Load model and tokenizer
tokenizer_16S = AutoTokenizer.from_pretrained('OceanOmics/eDNABERT-S_16S', trust_remote_code=True)
model_16S = AutoModel.from_pretrained('OceanOmics/eDNABERT-S_16S', trust_remote_code=True, config=config)
model_16S.to(device)
model_16S.eval()
names_16, seqs_16 = [], []
for seq in SeqIO.parse('16S_all_ASVs.fasta', 'fasta'):
    if 165 <= len(str(seq.seq)) <= 180:  # More efficient condition check
        seqs_16.append(str(seq.seq))
        names_16.append(str(seq.id))

print(f"Processing {len(seqs_16)} sequences")

batch_size = 32 # tested on an A100
num_sequences = len(seqs_16)
all_e_16 = np.zeros((num_sequences, 768))

with torch.no_grad():  # no gradient calculation for inference
    for i in tqdm(range(0, num_sequences, batch_size)):
        batch_seqs = seqs_16[i:i+batch_size]

        inputs = tokenizer_16S(batch_seqs, return_tensors='pt', padding=True)
        inputs = {k: v.to(device) for k, v in inputs.items()}  # Move inputs to device

        hidden_states = model_16S(**inputs)[0]

        for j, hidden_state in enumerate(hidden_states):
            embedding_mean = torch.mean(hidden_state, dim=0)
            all_e_16[i+j] = embedding_mean.cpu().numpy()  # Store directly in pre-allocated array

print("Running TSNE...")
X_embedded = TSNE(
    n_components=2,
    learning_rate='auto',
    init='random',
    perplexity=50,  # Reasonable value
    n_jobs=-1  # Use all available cores
).fit_transform(all_e_16)

print("Saving results...")
with open('odr_all_tsne_16S.optimized.tsv', 'w') as out:
    for a, name in zip(X_embedded, names_16):
        out.write('\t'.join(map(str, list(a) + name.split("XXX"))) + '\n')

python main.py [.. skipping data flags..]
  --seed 1
  --max_length 2000
  --train_batch_size 8 
  --val_batch_size 8
  --lr 3e-06
  --lr_scale 100
  --epochs 3
  --feat_dim 128
  --temperature 0.05
  --con_method same_species
  --mix
  --mix_alpha 1.0
  --mix_layer_num -1
  --curriculum