update model and config

README.md

@@ -84,16 +84,19 @@ user_message = prompt_template.format(input=prompt, response_a=responseA, respon
 
 conversation = [{"role": "user", "content": user_message}]
 
-device = "cuda:0"
-model_name = "Skywork/Skywork-Critic-Llama3.1-8B"
+model_name = "Skywork/Skywork-Critic-Llama3.1-70B"
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype="auto",
+    device_map="auto"
+)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map=device,)
 
 input_ids = tokenizer.apply_chat_template(
     conversation, 
     tokenize=True, 
-    add_generation_prompt=True, 
-    return_tensors="pt").to(device)
+    add_generation_prompt=True,
+    return_tensors="pt").to(model.device)
 
 generation = model.generate(
     input_ids=input_ids,
@@ -142,4 +145,4 @@ If you find our work helpful, please feel free to cite us using the following Bi
   howpublished={\url{https://huggingface.co/Skywork}},
   url={https://huggingface.co/Skywork},
 }
-```
+```