ra-tip

This is the repository for UIUC CSE 546 Final Project: Unlocking Chain-of-Thought Reasoning in Small-Scale LLMs via Pseudo-Code Aligned Fine-Tuning. We release our transformed dataset and trained model at huggingface.

Installation

conda create ra-tip-env python=3.9
conda activate ra-tip-env
pip install -r requirement.txt

Reasoning Analyst

To train a long CoT chunker using the training data that we distilled from DeepSeek V3.1, run:

  python training/train_chunker.py \
  --train_file reasoning_analyst/data/chunker_train.jsonl \
  --validation_file reasoning_analyst/data/chunker_test.jsonl \
  --instruction_file prompts/chunk.txt \
  --model_name_or_path meta-llama/Llama-3.1-8B-Instruct \
  --output_dir <chunker_model_path> \
  --label_list "Common Segment,Exploration,Subgoal Setting,Verification" \
  --per_device_train_batch_size 1 \
  --gradient_accumulation_steps 1 \
  --num_train_epochs 1 \
  --bf16 \
  --eval_steps 0 \
  --eval_generation_steps 0 \
  --max_length 4096 \
  --max_eval_length 4096 \
  --report_to wandb \
  --max_new_tokens 4096 \
  --wandb_project "reasoning_behavior_chunk_final"

To train a reasoning behavior classifier on the long CoTs chunks, run:

python training/train_classifier.py \
  --train_file reasoning_analyst/data/classifier_train.jsonl \
  --validation_file reasoning_analyst/data/classifier_test.jsonl \
  --instruction_file prompts/classify.txt  \
  --model_name_or_path meta-llama/Llama-3.1-8B-Instruct \
  --output_dir <classifier_model_path> \
  --label_list "Common Segment,Exploration,Subgoal Setting,Verification" \
  --per_device_train_batch_size 1 \
  --gradient_accumulation_steps 1 \
  --num_train_epochs 1 \
  --bf16 \
  --eval_steps 100 \
  --eval_generation_steps 500 \
  --max_length 4096 \
  --max_eval_length 4096 \
  --report_to wandb \
  --wandb_project "reasoning_behavior_classification_final"

To run the pipeline using these trained model on long CoTs in s1:

python training/inference_integrated.py \
  --input_file data/data_retrieval2/qwen-math-k=3_random.jsonl \
  --output_file <chunks_path> \
  --model_name_or_path_chunk  <chunker_model_path> \
  --model_name_or_path_classify <classifier_model_path> \
  --instruction_file_chunk prompts/chunk.txt \
  --instruction_file_classify prompts/classify.txt \
  --bf16

Think In Pseudo-code

To filter the constructed sequence of reasoning behaviors:

python think-in-code/filter_chunks.py <chunks_path> <filtered_chunks_path>
python think-in-code/pseudo_code_output_str.py --input <filtered_chunks_path> --output <hf_repos>

We use s1's repo for training the model, changing the base model to Qwen2.5-7B-Instruct and block_size=30000.

For evaluation, we use OLMES. Replace oe_eval/configs/tasks.py in their repo with olmes_patch/tasks.py in ours and run:

• GPQA Diamond

PYTHONPATH=. python olmes/oe_eval/run_eval.py \
	--task gpqa_diamond:0shot_cot::code \
	--model <trained model path> \
	--matching_key id \
	--save-raw-requests true \
	--ctx_key ctxs \
	--model-type vllm \
	--random-subsample-seed 2025 \
	--model-args '{"max_length": 32768}' \
	--output-dir <output_dir>

• MATH 500

PYTHONPATH=. python olmes/oe_eval/run_eval.py \
	--task minerva_math_500::code \
	--model <trained model path> \
	--matching_key id \
	--save-raw-requests true \
	--ctx_key ctxs \
	--model-type vllm \
	--random-subsample-seed 2025 \
	--model-args '{"max_length": 32768}' \
	--output-dir <output_dir>