Axiom Python 1.5B

Axiom Python 1.5B is a text generation (causal language model) fine-tuned on Qwen/Qwen2.5-1.5B with a focus on Python programming and code generation.

The model was trained using LoRA + SFT with the TRL library on the CodeAlpaca_20K and PythonCodeInstruct_18K datasets.

Model Details

Installation

Install the following packages to get started:

pip install transformers torch

If you are using a GPU, make sure you have installed a CUDA-compatible PyTorch version.

Usage

1. Using pipeline (Simplest Way)

from transformers import pipeline

generator = pipeline(
    "text-generation",
    model="coderian/axiom-python-1.5B",
    device_map="auto",
    torch_dtype="auto",
)

prompt = """### Instruction:
Write a Python function that reverses the elements of a list.

### Answer:
"""

output = generator(
    prompt,
    max_new_tokens=256,
    temperature=0.7,
    top_p=0.9,
    do_sample=True,
)

print(output[0]["generated_text"])

2. Using AutoModelForCausalLM

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "coderian/axiom-python-1.5B"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="auto",
)

model.eval()

prompt = """### Instruction:
Write a Python function that adds two numbers.

### Answer:
"""

inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=256,
        temperature=0.7,
        top_p=0.9,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id,
    )

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

print(response)

3. Using the Chat Template

Since the Qwen2.5 tokenizer supports the ChatML format, you can also use the model for chat-style conversations:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "coderian/axiom-python-1.5B"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="auto",
)

messages = [
    {"role": "system", "content": "You are Axiom, a helpful Python coding assistant."},
    {"role": "user", "content": "Write a Python function to check if a number is prime."},
]

text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
)

inputs = tokenizer(text, return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=256,
        temperature=0.7,
        top_p=0.9,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id,
    )

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

print(response)

Recommended Generation Parameters

Training Details

After training, the LoRA adapter was merged into the base model and released as a single file. You can also load the adapter directly using the peft library:

from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

base = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-1.5B",
    torch_dtype="auto",
    device_map="auto",
)

model = PeftModel.from_pretrained(base, "path/to/adapter")

Limitations

• It is a small 1.5B parameter model and may make mistakes on very complex and long code generation tasks.
• It was trained only on Python-focused datasets; performance in other languages is limited.
• The training data has a maximum length of 256 tokens; consistency may degrade in very long contexts.
• Generated code may not always be correct or safe. Review it before running.
• It may contain known limitations inherited from the training data regarding bias and harmful content.

Intended Usage Tips

• It performs best on single-line and medium-complexity Python functions.
• Lower the temperature value if you want stable output for code generation.
• Since the model was trained in a completion format, the ### Instruction: / ### Answer: template yields the highest quality output.
• For batched inference, remember to set tokenizer.pad_token = tokenizer.eos_token.

License

The base model Qwen2.5 is released under the Apache-2.0 license, and this model is also shared under the Apache-2.0 license.

Resources

• Base Model: Qwen/Qwen2.5-1.5B
• Training Library: TRL
• Dataset 1: HuggingFaceH4/CodeAlpaca_20K
• Dataset 2: iamtarun/python_code_instructions_18k_alpaca