from models import (
        HybridNoPEForCausalLM,
        HybridForCausalLM,
        HybridNoPEConfig,
        HybridConfig
    )

def get_model(args, tokenizer):
    if not args.nope:
        config = HybridConfig(
            layers=args.layers,
            bos_token_id=0,
            eos_token_id=0,
            hidden_size=args.hidden_size,
            intermediate_size=args.hidden_size*4,
            num_attention_heads=args.heads,
            d_model=args.hidden_size,
            ssm_cfg={"d_state": args.state_dim},
            vocab_size=len(tokenizer),
        )
    if args.nope:
        config = HybridNoPEConfig(
            layers=args.layers,
            bos_token_id=0,
            eos_token_id=0,
            hidden_size=args.hidden_size,
            intermediate_size=args.hidden_size*4,
            num_attention_heads=args.heads,
            d_model=args.hidden_size,
            ssm_cfg={"d_state": args.state_dim},
            vocab_size=len(tokenizer),
        )

    if not args.nope:
        model = HybridForCausalLM(config)
    if args.nope:
        model = HybridNoPEForCausalLM(config)
        
    return model