This model is trained to process Jupyter notebook content and generate structured outputs with think-tool-think chains. It produces interleaved reasoning and tool call outputs.

Base Model: Qwen/Qwen3-4B-Instruct-2507
Fine-tuning Method: LoRA (rank 8, alpha 16)
License: Apache 2.0

Usage

Installation

pip install transformers peft torch accelerate huggingface-hub tqdm bitsandbytes

Or use the provided requirements file:

pip install -r requirements_inference.txt

Load the Model

python

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch

# Load base model
base_model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-4B-Instruct-2507",
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True,
)

# Load LoRA adapter
model = PeftModel.from_pretrained(
    base_model,
    "AutomatedScientist/qwen-4b-jupyter-to-interleave-tool-call-lora"
)

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(
    "AutomatedScientist/qwen-4b-jupyter-to-interleave-tool-call-lora",
    trust_remote_code=True
)

Generate Output

python

# Prepare input
messages = [{"role": "user", "content": "Your Jupyter notebook content here..."}]
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(text, return_tensors="pt").to(model.device)

# Generate
with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=512,
        do_sample=True,
        temperature=0.7,
        pad_token_id=tokenizer.pad_token_id,
    )

# Decode
response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
print(response)

Run Inference Script

Download sample data from Aurora dataset:

python download_aurora_data.py --output-dir data/aurora_jupyter

Run inference:

bash

python inference.py \
    --adapter-path AutomatedScientist/qwen-4b-jupyter-to-interleave-tool-call-lora \
    --input-file data/aurora_jupyter/stage1/code/jupyter-structured-clean-dedup_part_aa.jsonl \
    --output-file results.jsonl \
    --max-samples 10

Sample Data

This model works with Jupyter notebook data. You can use the Aurora dataset for inference:

stage1/code/jupyter-structured-clean-dedup_part_aa.jsonl
stage1/code/jupyter-structured-clean-dedup_part_ab.jsonl

Training Details

Hyperparameters

Learning rate: 2e-4
Batch size: 1 (with gradient accumulation of 8)
Epochs: 3
LoRA rank: 8
LoRA alpha: 16
LoRA dropout: 0.05
Sequence length: 4096

Training Results

Training Loss	Epoch	Step	Validation Loss
0.1429	0.5	57	0.0904
0.057	1.0	114	0.0560
0.0531	1.5	171	0.0527
0.0506	2.0	228	0.0528
0.0511	2.5	285	0.0524

Final validation loss: 0.0524

Framework Versions

PEFT 0.18.0
Transformers 4.57.1
PyTorch 2.8.0+cu128
Datasets 4.4.1

Docker Inference (vLLM)

For high-throughput inference using vLLM with Docker (optimized for NVIDIA Blackwell GPUs):

Build the Docker image

docker build -t qwen-jupyter-inference .

Run inference

bash

docker run --gpus all --shm-size=8gb --ipc=host \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    -v $(pwd)/data:/data \
    qwen-jupyter-inference

Custom inference options

bash

docker run --gpus all --shm-size=8gb --ipc=host \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    -v $(pwd)/data:/data \
    qwen-jupyter-inference \
    python3 inference_vllm.py \
        --adapter-path AutomatedScientist/qwen-4b-jupyter-to-interleave-tool-call-lora \
        --input-file /data/your_input.jsonl \
        --output-file /data/output.jsonl \
        --max-samples 100 \
        --max-new-tokens 4096 \
        --max-model-len 32768

vLLM inference without Docker

bash

pip install vllm
python inference_vllm.py \
    --adapter-path AutomatedScientist/qwen-4b-jupyter-to-interleave-tool-call-lora \
    --input-file data/input.jsonl \
    --output-file data/output.jsonl \
    --max-samples 100

Evaluation

The model outputs are evaluated using a rubric that checks for:

Proper ## Instruction header
Proper ### Response header
Valid think/tool chain list
At least 3 chain entries

Run evaluation on outputs:

python

from rubric import get_reward, evaluate

text = "## Instruction\n...\n### Response\n..."
score = get_reward(text)  # Returns 0.0-1.0
details = evaluate(text)  # Returns detailed breakdown

Files

adapter_model.safetensors - LoRA adapter weights
adapter_config.json - LoRA configuration
tokenizer.json, tokenizer_config.json - Tokenizer files
inference.py - Basic inference script (transformers)
inference_vllm.py - High-throughput vLLM inference script
Dockerfile - Docker image for vLLM inference (Blackwell-optimized)
rubric.py - Evaluation rubric for output scoring
system_prompt.py - System prompt used during inference
download_aurora_data.py - Script to download sample data
requirements_inference.txt - Python dependencies

Citation

Built with Axolotl