| import os |
| from models import ( |
| LSTM, |
| GPTNeoXAlibiForCausalLM, |
| GPTNeoXHardAlibiForCausalLM, |
| GPTNeoXNoPEForCausalLM, |
| GPTNeoXForCausalLM, |
| MambaForCausalLM, |
| HybridNoPEForCausalLM, |
| HybridForCausalLM, |
| HybridNoPEConfig, |
| HybridConfig |
| ) |
|
|
| from transformers import GPTNeoXConfig |
| from transformers import MambaConfig |
| from transformers import AutoModelForCausalLM |
|
|
| def get_model(args, tokenizer): |
| if args.model in ["T_nope","T_rope","T_alibi"]: |
| config = GPTNeoXConfig( |
| bos_token_id=0, |
| eos_token_id=0, |
| hidden_size=args.hidden_size, |
| intermediate_size=args.hidden_size*4, |
| num_attention_heads=args.heads, |
| num_hidden_layers=args.layers, |
| vocab_size=len(tokenizer), |
| ) |
| elif args.model == "T_hard_alibi": |
| config = GPTNeoXConfig( |
| bos_token_id=0, |
| eos_token_id=0, |
| hidden_size=args.hidden_size, |
| intermediate_size=args.hidden_size*4, |
| num_attention_heads=args.heads, |
| num_hidden_layers=args.layers, |
| num_masked_heads=args.num_masked_heads, |
| vocab_size=len(tokenizer), |
| ) |
| elif args.model == "mamba": |
| config = MambaConfig( |
| hidden_size=args.hidden_size, |
| d_model=args.hidden_size, |
| n_layer=args.layers, |
| ssm_cfg={"d_state": args.state_dim}, |
| vocab_size=len(tokenizer), |
| ) |
| elif args.model == "hybrid": |
| config = HybridConfig( |
| bos_token_id=0, |
| eos_token_id=0, |
| hidden_size=args.hidden_size, |
| intermediate_size=args.hidden_size*4, |
| num_attention_heads=args.heads, |
| num_hidden_layers=args.layers, |
| d_model=args.hidden_size, |
| n_layer=args.layers, |
| ssm_cfg={"d_state": args.state_dim}, |
| vocab_size=len(tokenizer), |
| ) |
| elif args.model == "hybrid_nope": |
| config = HybridNoPEConfig( |
| bos_token_id=0, |
| eos_token_id=0, |
| hidden_size=args.hidden_size, |
| intermediate_size=args.hidden_size*4, |
| num_attention_heads=args.heads, |
| num_hidden_layers=args.layers, |
| d_model=args.hidden_size, |
| n_layer=args.layers, |
| ssm_cfg={"d_state": args.state_dim}, |
| vocab_size=len(tokenizer), |
| ) |
|
|
|
|
| |
| if args.model=="T_rope": |
| model = GPTNeoXForCausalLM(config) |
| elif args.model=="T_nope": |
| model = GPTNeoXNoPEForCausalLM(config) |
| elif args.model=="T_alibi": |
| model = GPTNeoXAlibiForCausalLM(config) |
| elif args.model=="T_hard_alibi": |
| model = GPTNeoXHardAlibiForCausalLM(config) |
| elif args.model=="mamba": |
| model = MambaForCausalLM(config) |
| elif args.model=="lstm": |
| model = LSTM( |
| embedding_dim=args.hidden_size, |
| vocab_size=len(tokenizer), |
| num_layers=args.layers, |
| dropout_rate=0.65 |
| ) |
| elif args.model=="hybrid": |
| model = HybridForCausalLM(config) |
| elif args.model=="hybrid_nope": |
| model = HybridNoPEForCausalLM(config) |
|
|
| if args.model == "pretrained": |
| model = AutoModelForCausalLM.from_pretrained(args.pretrain_model) |
| |
| return model |
|
|
|
|
|
|
