Update src/demo.py

src/demo.py

@@ -13,18 +13,17 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 TOKEN = os.environ.get("HF_TOKEN", None)
 
 
-model_id = "meta-llama/Llama-2-7b-chat-hf"
-# tokenizer = AutoTokenizer.from_pretrained(model_id)
-# model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16)
-
-
 type2dataset = {
-    "re2text-easy": load_dataset('3B-Group/ConvRe', "en-re2text", token=True, split="prompt1"),
-    "re2text-hard": load_dataset('3B-Group/ConvRe', "en-re2text", token=True, split="prompt4"),
-    "text2re-easy": load_dataset('3B-Group/ConvRe', "en-text2re", token=True, split="prompt1"),
-    "text2re-hard": load_dataset('3B-Group/ConvRe', "en-text2re", token=True, split="prompt3")
+    "re2text-easy": load_dataset('3B-Group/ConvRe', "en-re2text", token=TOKEN, split="prompt1"),
+    "re2text-hard": load_dataset('3B-Group/ConvRe', "en-re2text", token=TOKEN, split="prompt4"),
+    "text2re-easy": load_dataset('3B-Group/ConvRe', "en-text2re", token=TOKEN, split="prompt1"),
+    "text2re-hard": load_dataset('3B-Group/ConvRe', "en-text2re", token=TOKEN, split="prompt3")
 }
 
+model_id = "meta-llama/Llama-2-7b-chat-hf"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16)
+
 # type2dataset = {}