Trayvon commited on
Commit
789bfa3
·
1 Parent(s): c8073f1

Update src/demo.py

Browse files
Files changed (1) hide show
  1. src/demo.py +3 -3
src/demo.py CHANGED
@@ -20,9 +20,9 @@ type2dataset = {
20
  "text2re-hard": load_dataset('3B-Group/ConvRe', "en-text2re", token=TOKEN, split="prompt3")
21
  }
22
 
23
- model_id = "TheBloke/Llama-2-7b-Chat-GPTQ"
24
- tokenizer = AutoTokenizer.from_pretrained(model_id, token=TOKEN, use_fast=True)
25
- model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, token=TOKEN, trust_remote_code=False, revision="main").eval()
26
 
27
  # type2dataset = {}
28
 
 
20
  "text2re-hard": load_dataset('3B-Group/ConvRe', "en-text2re", token=TOKEN, split="prompt3")
21
  }
22
 
23
+ model_id = "meta-llama/Llama-2-7b-chat-hf"
24
+ tokenizer = AutoTokenizer.from_pretrained(model_id, token=TOKEN)
25
+ model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, token=TOKEN).eval()
26
 
27
  # type2dataset = {}
28