File size: 432 Bytes
4c13c46
 
 
2be1c87
 
4c13c46
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
from transformers import PretrainedConfig


class MaliciousCodeTestConfig(PretrainedConfig):
    model_type = "malicious-code-test"

    def __init__(self, vocab_size=32000, n_embd=128, n_layer=2, n_head=2, block_size=16, **kwargs):

        self.vocab_size = vocab_size
        self.n_embd = n_embd
        self.n_layer = n_layer
        self.n_head = n_head
        self.block_size = block_size
        super().__init__(**kwargs)