Update README.md

README.md

@@ -55,7 +55,7 @@ parser.add_argument("--resume_from_checkpoint", action="store_true", default=Fal
 parser.add_argument("--lora", action="store_true")
 args = parser.parse_args()
 
-qwq_dataset = load_dataset("amphora/QwQ-LongCoT-130K", split = "train")
+qwq_dataset = load_dataset("amphora/QwQ-LongCoT-130K-2", split = "train")
 messages = []
 for each in qwq_dataset:
     msg = [
@@ -69,10 +69,10 @@ TRAIN_SPLIT_RATIO = 0.9
 train_size = int(TRAIN_SPLIT_RATIO * len(messages))
 eval_size = len(messages) - train_size
 
-tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
+tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
 
 # The model to optimise
-model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct", torch_dtype=torch.bfloat16, device_map="auto") 
+model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct", torch_dtype=torch.bfloat16, device_map="auto")