Qwen2.5-Coder-1.5B - Cellsistant Tool Calling

Fine-tuned version of Qwen2.5-Coder-1.5B for tool calling in Cellsistant (AI assistant for JupyterLab).

Model Details

• Base Model: Qwen/Qwen2.5-Coder-1.5B-Instruct
• Training Framework: Unsloth (2x faster)
• Training Data: 2,342 examples for tool calling
• Training Time: 6.57 minutes on Tesla T4
• Training Loss: 0.39 → 0.00042

Supported Tools

The model supports 19 tools for Jupyter notebook manipulation:

Notebook Tools

• create_cell - Create new cells
• update_cell - Update cell content
• delete_cell - Delete cells
• get_cell_content - Read cell content
• get_notebook_content - Get full notebook
• execute_cell - Execute code cells
• get_cell_output - Get cell output
• analyze_image - Analyze plots/images

File Tools

• read_file - Read files
• write_file - Write files
• list_directory - List directories
• create_notebook - Create notebooks
• delete_file - Delete files
• rename_file - Rename/move files

Other Tools

• run_shell - Execute shell commands
• find_in_notebook - Search in notebook
• replace_in_cell - Replace text in cells
• install_package - Install Python packages
• list_packages - List installed packages

Usage

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# Load base model
base_model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-Coder-1.5B-Instruct",
    device_map="cuda",
    torch_dtype=torch.float16,
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-1.5B-Instruct")

# Load LoRA adapters
model = PeftModel.from_pretrained(base_model, "coconut495/qwen2.5-coder-1.5b-cellsistant-tool-calling")

# Generate tool calls
prompt = '''<|im_start|>system
You are a helpful assistant with access to the following tools:
[{"type": "function", "function": {"name": "create_cell", ...}}]
<|im_end|>
<|im_start|>user
Add a code cell that prints hello<|im_end|>
<|im_start|>assistant
'''

inputs = tokenizer([prompt], return_tensors="pt").to("cuda")
outputs = model.generate(**inputs, max_new_tokens=64, use_cache=False)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Training Details

Dataset

Training dataset contains 2,603 examples covering all 19 tools with various phrasings and edge cases.

License

Apache 2.0

Acknowledgements

• Unsloth for fast training
• Qwen Team for the base model