Update README.md

README.md

@@ -6,52 +6,15 @@ I do not own, distribute, or take credits for this model, all copyrights belong
 
 
 # how to load the model 
-download the weights
-```
-!git clone https://huggingface.co/not-lain/TunBERT
-```
-load the model
 ```python
-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)
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-    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")
+tokenizer = AutoTokenizer.from_pretrained("not-lain/TunBERT")
+model = AutoModelForSequenceClassification.from_pretrained("not-lain/TunBERT",trust_remote_code=True)
 ```
 
+**IMPORTANT** : make sure to enable `trust_remote_code=True` 
+
 
 # how to use the model 
 ```python