File size: 879 Bytes
8a00d0d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 |
from transformers import AutoTokenizer, ModernBertModel
import torch
import torch.nn as nn
import torch.optim as optim
import pdb
class modernBERT(nn.Module):
def __init__(self, model_name="answerdotai/ModernBERT-base"):
super(modernBERT, self).__init__()
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.bert = ModernBertModel.from_pretrained(model_name)
def forward(self, inputs):
# inputs = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True)
outputs = self.bert(**inputs)
return outputs.last_hidden_state # logits
# Example training loop
if __name__ == "__main__":
model = modernBERT("answerdotai/ModernBERT-base")
texts = ["Potato's no name for a dog"]
text_inputs = {"input_ids": model.tokenizer(texts)}
output = model(text_inputs)
print(output[0].shape)
|