Update README.md

README.md

@@ -29,15 +29,14 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 
 
 
-tokenizer = AutoTokenizer.from_pretrained("SimpleBerry/LLaMA-O1-Supervised-1129")
-model = AutoModelForCausalLM.from_pretrained("SimpleBerry/LLaMA-O1-Supervised-1129",device_map='auto')
+tokenizer = AutoTokenizer.from_pretrained("/mnt/hwfile/ai4chem/CKPT/longcot_sft_llama3.1_ZD_11_29_1/")
+model = AutoModelForCausalLM.from_pretrained("/mnt/hwfile/ai4chem/CKPT/longcot_sft_llama3.1_ZD_11_29_1/",device_map='auto')
 
 
 
 template = "<start_of_father_id>-1<end_of_father_id><start_of_local_id>0<end_of_local_id><start_of_thought><problem>{content}<end_of_thought><start_of_rating><positive_rating><end_of_rating>\n<start_of_father_id>0<end_of_father_id><start_of_local_id>1<end_of_local_id><start_of_thought><expansion>"
 
-def llama_o1_template(data):
-    query = data['query']
+def llama_o1_template(query):
     text = template.format(content=query)
     return text
 
@@ -51,7 +50,6 @@ def batch_predict(input_texts):
     assitant_responses = [item for i,item in enumerate(response_texts)]
     return assitant_responses
 
-
 i = 'If Diana needs to bike 10 miles to reach home and she can bike at a speed of 3 mph for two hours before getting tired, and then at a speed of 1 mph until she reaches home, how long will it take her to get home?'
 input_texts = [llama_o1_template(i)]
 assitant_responses = batch_predict(input_texts)