from transformers import LlamaConfig

class GRetrieverConfig(LlamaConfig):
    model_type = "llama"
    
    def __init__(
        self,
        max_txt_len: int = 1024,
        max_new_tokens: int = 256,
        gnn_num_layers: int = 4,
        gnn_in_dim: int = 768,
        gnn_hidden_dim: int = 1024,
        gnn_num_heads: int = 4,
        gnn_dropout: int = 0,
        bos_id: list = [128000, 128000, 128006, 882, 128007],
        **kwargs
    ):
        pretrained_config = LlamaConfig.from_pretrained("NousResearch/Hermes-3-Llama-3.1-8B")
        pretrained_config.update(kwargs)
        
        self.max_txt_len = max_txt_len
        self.max_new_tokens = max_new_tokens
        self.gnn_num_layers = gnn_num_layers
        self.gnn_in_dim = gnn_in_dim
        self.gnn_hidden_dim = gnn_hidden_dim
        self.gnn_num_heads = gnn_num_heads
        self.gnn_dropout = gnn_dropout
        self.bos_id = bos_id
        
        super().__init__(**pretrained_config.to_dict())
        self.pad_token_id = pretrained_config.eos_token_id