teknium commited on
Commit
395701f
1 Parent(s): 7da7c54

Update transformers_inference_example.py

Browse files
Files changed (1) hide show
  1. transformers_inference_example.py +2 -2
transformers_inference_example.py CHANGED
@@ -3,11 +3,11 @@
3
 
4
  import torch
5
  from transformers import AutoTokenizer, AutoModelForCausalLM
6
- from transformers import LlamaTokenizer, LlamaForCausalLM, MistralForCausalLM
7
  import bitsandbytes, flash_attn
8
 
9
  tokenizer = LlamaTokenizer.from_pretrained('NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO', trust_remote_code=True)
10
- model = MistralForCausalLM.from_pretrained(
11
  "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
12
  torch_dtype=torch.float16,
13
  device_map="auto",
 
3
 
4
  import torch
5
  from transformers import AutoTokenizer, AutoModelForCausalLM
6
+ from transformers import LlamaTokenizer, MixtralForCausalLM
7
  import bitsandbytes, flash_attn
8
 
9
  tokenizer = LlamaTokenizer.from_pretrained('NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO', trust_remote_code=True)
10
+ model = MixtralForCausalLM.from_pretrained(
11
  "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
12
  torch_dtype=torch.float16,
13
  device_map="auto",