S3-CoT-Llama-3.1-8B-Instruct

We release trained checkpoints in our paper (S3-CoT: Self-Sampled Succinct Reasoning Enables Efficient Chain-of-Thought LLMs).

• Model Checkpoints:

• Inference Template: we provide System 1 and System 2 templates used in our training. During inference, please use the System 1 template to reproduce the results reported in our paper.

from transformers import AutoModelForCausalLM, AutoTokenizer

system1_template = " Please provide as a brief reasoning process as possible, and put your final answer within \\boxed{}"
system2_template = " Please reason step by step, and put your final answer within \\boxed{}"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype="bfloat16",
)
generation_config={"xxx"}
prompt = "xxx"

messages = [
    {"role": "user", "content": prompt+system1_template}
]
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)

model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
generated_ids = model.generate(
    **model_inputs,
    **generation_config,
)

output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
output_content = tokenizer.decode(output_ids, skip_special_tokens=True)