| 71 | |
| 72 | |
| 73 | def load_hf_model(model_name: str, device:str = 'cuda') -> nn.Module: |
| 74 | if "mamba" in model_name: |
| 75 | |
| 76 | # SA: can't pass in device here https://github.com/pytorch/pytorch/issues/10622 |
| 77 | model = MambaLMHeadModel.from_pretrained(model_name, device=device, dtype=torch.float16) |
| 78 | |
| 79 | else: |
| 80 | if "Mixtral" in model_name: |
| 81 | model = AutoModelForCausalLM.from_pretrained( |
| 82 | model_name, trust_remote_code=True, use_flash_attention_2=True, |
| 83 | # load_in_8bit=True, |
| 84 | torch_dtype=torch.bfloat16, device_map="auto" |
| 85 | ) |
| 86 | else: |
| 87 | model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, token="your token here") |
| 88 | model.to(device) |
| 89 | |
| 90 | |
| 91 | try: model.device = device |
| 92 | except: pass |
| 93 | model.eval() |
| 94 | return model |
| 95 | |
| 96 | |
| 97 | def load_tokenizer(model_name: str, is_hf: bool=False) -> nn.Module: |