TunBERT / README.md
not-lain's picture
Update README.md
0526e55
|
raw
history blame
2.17 kB
metadata
license: mit

Disclamer

I do not own, distribute, or take credits for this model, all copyrights belong to Instadeep under the MIT licence

how to load the model

download the weights

!git clone https://huggingface.co/not-lain/TunBERT

load the model

import torch.nn as nn
import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM, AutoModelForSequenceClassification, PreTrainedModel,AutoConfig, BertModel
from transformers.modeling_outputs import SequenceClassifierOutput
config = AutoConfig.from_pretrained("not-lain/TunBERT")
class classifier(nn.Module):
  def __init__(self,config):
    super().__init__()

    self.layer0 = nn.Linear(in_features=config.hidden_size, out_features=config.hidden_size, bias=True)
    self.layer1 = nn.Linear(in_features=config.hidden_size, out_features=config.type_vocab_size, bias=True)
  def forward(self,tensor):
    out1 = self.layer0(tensor)
    return self.layer1(out1)


class TunBERT(PreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.BertModel = BertModel(config)
        self.dropout = nn.Dropout(p=0.1, inplace=False)
        self.classifier = classifier(config)

    def forward(self,input_ids=None,token_type_ids=None,attention_mask=None,labels=None) :
      outputs = self.BertModel(input_ids,token_type_ids,attention_mask)
      sequence_output = self.dropout(outputs.last_hidden_state)
      logits = self.classifier(sequence_output)
      loss =None
      if labels is not None :
        loss_func = nn.CrossentropyLoss()
        loss = loss_func(logits.view(-1,self.config.type_vocab_size),labels.view(-1))
      return SequenceClassifierOutput(loss = loss, logits= logits, hidden_states=outputs.last_hidden_state,attentions=outputs.attentions)


tunbert = TunBERT(config)
tunbert.load_state_dict(torch.load("TunBERT/pytorch_model.bin"))

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

how to use the model

text = "[insert text here]"
inputs = tokenizer(text,return_tensors='pt')
output = model(**inputs)