Update src/demo.py

src/demo.py

@@ -22,7 +22,7 @@ type2dataset = {
 
 model_id = "meta-llama/Llama-2-7b-chat-hf"
 tokenizer = AutoTokenizer.from_pretrained(model_id, token=TOKEN)
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, token=TOKEN)
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, token=TOKEN)
 
 # type2dataset = {}
 
@@ -35,7 +35,7 @@ def generate(input_text, sys_prompt) -> str:
 '''
     input_str = sys_prompt + input_text + " [/INST]"
     
-    input_ids = tokenizer(input_str, return_tensors="pt").input_ids
+    input_ids = tokenizer(input_str, return_tensors="pt", torch_dtype=torch.bfloat16).input_ids
     outputs = model.generate(input_ids, max_length=512)
 
     result = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]