Update README.md

README.md

@@ -7,7 +7,7 @@ tags:
 - full
 - generated_from_trainer
 model-index:
-- name: longcot_sft_llama3.1_ZD_11_29_1
+- name: SimpleBerry/LLaMA-O1-Supervised-1129
   results: []
 ---
 
@@ -16,45 +16,45 @@ should probably proofread and complete it, then remove this comment. -->
 
 # longcot_sft_llama3.1_ZD_11_29_1
 
-This model is a fine-tuned version of [/mnt/hwfile/ai4chem/CKPT/longcot_pt_llama3.1_ZD_11_29_1/](https://huggingface.co//mnt/hwfile/ai4chem/CKPT/longcot_pt_llama3.1_ZD_11_29_1/) on the longcot_sft_1 dataset.
+This model is a fine-tuned version of [SimpleBerry/LLaMA-O1-Base-1127](https://huggingface.co/SimpleBerry/LLaMA-O1-Base-1127) on the [SimpleBerry/OpenLongCoT-SFT](SimpleBerry/OpenLongCoT-SFT) dataset.
 
-## Model description
+# Inference
+```Python
+import json
+import datasets
+import torch
+import random
+import numpy as np
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-More information needed
 
-## Intended uses & limitations
 
-More information needed
+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')
 
-## Training and evaluation data
 
-More information needed
 
-## Training procedure
+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>"
 
-### Training hyperparameters
+def llama_o1_template(data):
+    query = data['query']
+    text = template.format(content=query)
+    return text
 
-The following hyperparameters were used during training:
-- learning_rate: 5e-06
-- train_batch_size: 1
-- eval_batch_size: 8
-- seed: 42
-- distributed_type: multi-GPU
-- num_devices: 24
-- gradient_accumulation_steps: 16
-- total_train_batch_size: 384
-- total_eval_batch_size: 192
-- optimizer: Use OptimizerNames.ADAMW_TORCH with betas=(0.9,0.999) and epsilon=1e-08 and optimizer_args=No additional optimizer arguments
-- lr_scheduler_type: cosine
-- num_epochs: 2.0
 
-### Training results
+def batch_predict(input_texts):
+    input_texts = [input_text.replace('<|end_of_text|>','') for input_text in input_texts]
+    inputs = tokenizer(input_texts, return_tensors="pt").to(model.device)
+    responses = model.generate(**inputs, max_new_tokens=1024)
+    response_texts = tokenizer.batch_decode(responses, skip_special_tokens=False)
+    # assitant_responses = [item[len(input_texts[i]):] for i,item in enumerate(response_texts)]
+    assitant_responses = [item for i,item in enumerate(response_texts)]
+    return assitant_responses
 
 
+i = input()
+input_texts = [llama_o1_template(i)]
+assitant_responses = batch_predict(input_texts)
+print(assitant_responses)
 
-### Framework versions
-
-- Transformers 4.46.2
-- Pytorch 2.3.1
-- Datasets 3.1.0
-- Tokenizers 0.20.1
+```