BERT_2_Labels / README.md
AkashKhamkar's picture
Adding model details.
e7ed86c
|
raw
history blame
772 Bytes
metadata
license: apache-2.0

Number of Epochs = 5
Dataset Size = 5.5 k samples [train/validation]
Number of labels used = 2
Thresholding = True
Thresholding value = 0.7

Below is the function to aplly thresholding to output logits.

  def get_prediction(text):
    encoding = new_tokenizer(text, return_tensors="pt", padding="max_length", truncation=True, max_length=128)
    encoding = {k: v.to(trainer.model.device) for k,v in encoding.items()}

    outputs = new_model(**encoding)

    logits = outputs.logits

    sigmoid = torch.nn.Sigmoid()
    probs = probs.detach().numpy()
    label = np.argmax(probs, axis=-1)
    if label == 1:
      if probs[1] > 0.7:
        return 1
      else:
        return 0
    else:
      return 0