INNER CODE UNIT · Python
BertClassification
yya518/FinBERT · archive/bertModel.py:7
class BertClassification(nn.Module):
def __init__(self, weight_path, num_labels=2, vocab="base-cased"):
super(BertClassification, self).__init__()
self.num_labels = num_labels
self.vocab = vocab
if self.vocab == "base-cased":
self.bert = BertModel.from_pretrained(weight_path)
self.config = BertConfig(vocab_size_or_config_json_file=28996, hidden_size=768, num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)
elif self.vocab == "base-uncased":
self.bert = BertModel.from_pretrained(weight_path)
self.config = BertConfig(vocab_size_or_config_json_file=30522, hidden_size=768, num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)
elif self.vocab == "finance-cased":
self.bert = BertModel.from_pretrained(weight_path)
self.config = BertConfig(vocab_size_or_config_json_file=28573, hidden_size=768, num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)