qwen35-2b-serax-codes-base API & Inference Endpoint

Model description

Base: Qwen/Qwen3.5-2B (unified VLM class, used text-only).
Architecture: Qwen3_5ForConditionalGeneration, Gated DeltaNet hybrid attention, 24 layers, hidden 2048, vocab 248,320, tie_word_embeddings.
Stage-1 adapter (now merged): LoRA r=64, alpha=128, dropout 0.0, bf16, targeting q/k/v/o_proj, gate/down/up_proj, and the Qwen3.5 linear-attention projections (linear_attn.in_proj_qkv, linear_attn.in_proj_z, linear_attn.out_proj).
What it knows: forward/reverse lookup, description, and child-enumeration of UNSPSC / HS / NAICS / SOC / ISO-3166 codes, and their SERAX encodings.

Intended use

Use as the base model for SERAX task adapters (e.g. the v1.5 business-website analyst). It can also be queried directly for code lookups (see below), but it is a building block, not a finished product model.

How to use

Loads as an ordinary causal-LM. There is no system prompt for this stage — training rows are direct user → assistant code queries.

python
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

m = "vantigeai/qwen35-2b-serax-codes-base"
tok = AutoTokenizer.from_pretrained(m)
model = AutoModelForCausalLM.from_pretrained(m, torch_dtype=torch.bfloat16, device_map="auto")

msgs = [{"role": "user",
         "content": "What UNSPSC v26.0801 codes belong directly to 31163000 — Couplings?"}]
inputs = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt").to(model.device)
out = model.generate(inputs, max_new_tokens=256, do_sample=False)
print(tok.decode(out[0, inputs.shape[1]:], skip_special_tokens=True))

To load the stage-2 analyst adapter on top of this base:

python
from peft import PeftModel
model = PeftModel.from_pretrained(model, "vantigeai/serax_business_website_analyst_v1.5")

Training data

Stage-1 data is data/code_train.jsonl — 31,817 examples, each a direct user → assistant turn (no system prompt), generated from the authoritative code tables (prepare_code_data.py from data/code_fine_tuning.parquet). Only the assistant turn is trained on (roles_to_train: ["assistant"]).

Coverage by code system (classified by the query text):

Table with columns: Code system, Examples
Code system	Examples
UNSPSC v26.0801	12,601
HS 2022	5,326
NAICS 2022	5,110
SOC 2018	4,580
Other (ISO-3166 geo, SERAX code records, etc.)	4,200
Total	31,817

Task shapes (examples):

markdown
U: Give the description for SOC 2018 code 13-2030.
A: Budget Analysts

U: Find the SOC 2018 code for Physical Scientists, All Other.
A: 19-2099

U: What UNSPSC v26.0801 codes belong directly to 31163000 — Couplings?
A: 31163001 Elastomeric couplings
   31163002 Gear couplings
   31163003 Metallic couplings
   ...

U: Emit the SERAX record for ISO-3166 code KG-B.
A: ⨠ISO-3166-2⨡KG-B⨢Batken

U: Emit the SERAX record for Wood; assembled flooring panels … of bamboo in HS 2022.
A: ⨥441873⨦Wood; assembled flooring panels, of bamboo or with at least the top layer …

The tasks cover: code→description, description→code, enumerate children under a parent code, and emit the SERAX encoding of a code. Assistant length: min 2 / median 54 / max 12,438 characters.

Training configuration

Trained with Axolotl as stage 1 of 2. Key choices vs. stage 2: higher LR (2e-4) to drive rote memorization, short sequence_len (3072 — the dataset max is 105 in + 2863 out tokens), and a large micro_batch_size (32) since rows are short and pack efficiently.

yaml
base_model: Qwen/Qwen3.5-2B
# Stage 1 of 2: KNOWLEDGE INJECTION.
# Teach the model the SERAX / UNSPSC v26.0801 / HS-2022 code vocabulary for the
# exact version we use, BEFORE the stage-2 task fine-tune (axolotl_qwen_3.5_2b.yaml).
# Qwen3.5 loads as a unified VLM class; this is a TEXT-ONLY fine-tune.
plugins:
  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin   # essential for the 248k vocab
strict: false
chat_template: qwen3_5
datasets:
  # Produced by prepare_code_data.py from data/code_fine_tuning.parquet.
  # Same messages/qwen3_5 format as stage 2, so the model learns the codes
  # under the exact chat template it will be evaluated with.
  - path: ./data/code_train.jsonl
    type: chat_template
    field_messages: messages
    message_property_mappings:
      role: role
      content: content
    roles_to_train: ["assistant"]
dataset_prepared_path: last_run_prepared_codes
val_set_size: 0.02
output_dir: ./outputs/qwen35-2b-lora-codes
# ---- sequence / packing ----
# Dataset max is 105 in + 2863 out = 2968 tokens; 3072 holds every row plus
# chat-template overhead with no truncation. Median is ~50 tokens, so packing
# is extremely effective here.
sequence_len: 3072
sample_packing: true
eval_sample_packing: true
pad_to_sequence_len: true      # keeps packed-batch VRAM flat (OOM-safety)
# ---- adapter: LoRA (bf16) ----
# Higher rank than stage 2: rote vocabulary memorization benefits from more
# adapter capacity. The resulting adapter is the starting point for stage 2.
adapter: lora
lora_r: 64
lora_alpha: 128
lora_dropout: 0.0
lora_target_modules:
  - q_proj
  - k_proj
  - v_proj
  - o_proj
  - gate_proj
  - down_proj
  - up_proj
  # Qwen3.5-specific Gated DeltaNet linear-attention projections:
  - linear_attn.in_proj_qkv
  - linear_attn.in_proj_z
  - linear_attn.out_proj
# ---- attention / precision ----
attn_implementation: flash_attention_2
bf16: auto
tf32: true
# ---- optimizer / schedule ----
optimizer: adamw_torch        # or adamw_torch_8bit to save a little VRAM
lr_scheduler: cosine
learning_rate: 2.0e-4         # a touch higher than stage 2 to drive memorization
weight_decay: 0.01
warmup_ratio: 0.03
num_epochs: 3
micro_batch_size: 32          # A100 80GB is barely loaded at 8; do the work in one efficient pass
gradient_accumulation_steps: 1   # effective batch 32 (same as before, but no sequential micro-steps)
# ---- dataloader (keep the GPU fed; default num_workers=0 starves it) ----
dataloader_num_workers: 8        # 12 CPU cores -> 8 leaves headroom for the main proc
dataloader_pin_memory: true
dataloader_prefetch_factor: 4
# ---- memory ----
gradient_checkpointing: true
gradient_checkpointing_kwargs:
  use_reentrant: false
# ---- logging / saving ----
logging_steps: 1
evals_per_epoch: 2
saves_per_epoch: 1
save_total_limit: 3
# ---- monitoring (Weights & Biases) ----
wandb_project: qwen35-2b-lora
wandb_entity:                  # your username or team; blank = default
wandb_name: qwen35-2b-lora-codes-4k-r64
wandb_watch:                   # leave blank; "gradients"/"all" adds overhead
wandb_log_model:               # "checkpoint" to upload adapters, blank to skip

Key hyperparameters

Table

adapter	LoRA r=64, alpha=128, dropout 0.0 (merged post-training)
learning_rate	2e-4 (cosine, warmup_ratio 0.03)
epochs	3
micro_batch_size	32
gradient_accumulation_steps	1 (effective batch 32)
optimizer	adamw_torch
weight_decay	0.01
sequence_len	3,072 (sample packing)
precision	bf16, tf32

License

Apache-2.0, inheriting the Qwen3.5-2B license. This repository redistributes Qwen3.5-2B weights with a merged fine-tuning delta.

python

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

m = "vantigeai/qwen35-2b-serax-codes-base"
tok = AutoTokenizer.from_pretrained(m)
model = AutoModelForCausalLM.from_pretrained(m, torch_dtype=torch.bfloat16, device_map="auto")

msgs = [{"role": "user",
         "content": "What UNSPSC v26.0801 codes belong directly to 31163000 — Couplings?"}]
inputs = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt").to(model.device)
out = model.generate(inputs, max_new_tokens=256, do_sample=False)
print(tok.decode(out[0, inputs.shape[1]:], skip_special_tokens=True))

Code system

Examples

UNSPSC v26.0801

12,601

HS 2022

5,326

NAICS 2022

5,110

SOC 2018

4,580

Other (ISO-3166 geo, SERAX code records, etc.)

4,200

Total

31,817

markdown

U: Give the description for SOC 2018 code 13-2030.
A: Budget Analysts

U: Find the SOC 2018 code for Physical Scientists, All Other.
A: 19-2099

U: What UNSPSC v26.0801 codes belong directly to 31163000 — Couplings?
A: 31163001 Elastomeric couplings
   31163002 Gear couplings
   31163003 Metallic couplings
   ...

U: Emit the SERAX record for ISO-3166 code KG-B.
A: ⨠ISO-3166-2⨡KG-B⨢Batken

U: Emit the SERAX record for Wood; assembled flooring panels … of bamboo in HS 2022.
A: ⨥441873⨦Wood; assembled flooring panels, of bamboo or with at least the top layer …

yaml

base_model: Qwen/Qwen3.5-2B
# Stage 1 of 2: KNOWLEDGE INJECTION.
# Teach the model the SERAX / UNSPSC v26.0801 / HS-2022 code vocabulary for the
# exact version we use, BEFORE the stage-2 task fine-tune (axolotl_qwen_3.5_2b.yaml).
# Qwen3.5 loads as a unified VLM class; this is a TEXT-ONLY fine-tune.
plugins:
  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin   # essential for the 248k vocab
strict: false
chat_template: qwen3_5
datasets:
  # Produced by prepare_code_data.py from data/code_fine_tuning.parquet.
  # Same messages/qwen3_5 format as stage 2, so the model learns the codes
  # under the exact chat template it will be evaluated with.
  - path: ./data/code_train.jsonl
    type: chat_template
    field_messages: messages
    message_property_mappings:
      role: role
      content: content
    roles_to_train: ["assistant"]
dataset_prepared_path: last_run_prepared_codes
val_set_size: 0.02
output_dir: ./outputs/qwen35-2b-lora-codes
# ---- sequence / packing ----
# Dataset max is 105 in + 2863 out = 2968 tokens; 3072 holds every row plus
# chat-template overhead with no truncation. Median is ~50 tokens, so packing
# is extremely effective here.
sequence_len: 3072
sample_packing: true
eval_sample_packing: true
pad_to_sequence_len: true      # keeps packed-batch VRAM flat (OOM-safety)
# ---- adapter: LoRA (bf16) ----
# Higher rank than stage 2: rote vocabulary memorization benefits from more
# adapter capacity. The resulting adapter is the starting point for stage 2.
adapter: lora
lora_r: 64
lora_alpha: 128
lora_dropout: 0.0
lora_target_modules:
  - q_proj
  - k_proj
  - v_proj
  - o_proj
  - gate_proj
  - down_proj
  - up_proj
  # Qwen3.5-specific Gated DeltaNet linear-attention projections:
  - linear_attn.in_proj_qkv
  - linear_attn.in_proj_z
  - linear_attn.out_proj
# ---- attention / precision ----
attn_implementation: flash_attention_2
bf16: auto
tf32: true
# ---- optimizer / schedule ----
optimizer: adamw_torch        # or adamw_torch_8bit to save a little VRAM
lr_scheduler: cosine
learning_rate: 2.0e-4         # a touch higher than stage 2 to drive memorization
weight_decay: 0.01
warmup_ratio: 0.03
num_epochs: 3
micro_batch_size: 32          # A100 80GB is barely loaded at 8; do the work in one efficient pass
gradient_accumulation_steps: 1   # effective batch 32 (same as before, but no sequential micro-steps)
# ---- dataloader (keep the GPU fed; default num_workers=0 starves it) ----
dataloader_num_workers: 8        # 12 CPU cores -> 8 leaves headroom for the main proc
dataloader_pin_memory: true
dataloader_prefetch_factor: 4
# ---- memory ----
gradient_checkpointing: true
gradient_checkpointing_kwargs:
  use_reentrant: false
# ---- logging / saving ----
logging_steps: 1
evals_per_epoch: 2
saves_per_epoch: 1
save_total_limit: 3
# ---- monitoring (Weights & Biases) ----
wandb_project: qwen35-2b-lora
wandb_entity:                  # your username or team; blank = default
wandb_name: qwen35-2b-lora-codes-4k-r64
wandb_watch:                   # leave blank; "gradients"/"all" adds overhead
wandb_log_model:               # "checkpoint" to upload adapters, blank to skip

adapter

LoRA r=64, alpha=128, dropout 0.0 (merged post-training)

learning_rate

2e-4 (cosine, warmup_ratio 0.03)

epochs

micro_batch_size

gradient_accumulation_steps

1 (effective batch 32)

optimizer

adamw_torch

weight_decay

0.01

sequence_len

3,072 (sample packing)

precision

bf16, tf32

qwen35-2b-serax-codes-base

Get help setting up a custom Dedicated Endpoints.

README

Model description

Intended use

How to use

Training data

Training configuration

Key hyperparameters

License

Explore FriendliAI today

README

Model description

Intended use

How to use

Training data

Training configuration

Key hyperparameters

License