omartra999/tlm-xs-refactor

TLM-XS: a tiny Python refactoring model (3.13M parameters)

TLM-XS is a small encoder-decoder transformer trained from scratch, with no pretrained weights. Given a Python function and a task, it produces the edit as a SEARCH/REPLACE block that can be applied to the source.

Tasks: rename, simplify, deduplicate, remove_dead_code, extract_method.

• Code, training pipeline and evaluation: https://github.com/mutabalShawandar/tlm-code-refactoring
• Sizes in this series: TLM-XS, TLM-S, TLM-M

Architecture

4+4 layers, d_model 128, d_ff 512, 4 heads, 10k BPE vocabulary. Pre-Norm, rotary position embeddings on self-attention, tied embeddings, a copy-pointer head that blends copying from the input with generating, and per-layer task embeddings on the encoder.

Training

Span-corruption pretraining on a 5-tier curriculum, a bridging stage into SEARCH/REPLACE format, then supervised fine-tuning (25 epochs). This checkpoint is the SFT result, with no RL step. Final validation NLL 0.3170.

Results

Held-out set, 750 examples per task. "applied_ok" means the edit applies cleanly and passes the task-specific check.

rename	simplify	deduplicate	remove_dead_code	extract_method	overall
25.3%	30.5%	14.1%	75.2%	16.4%	32.3%

For reference, CodeT5-small (60M parameters, official pretrained weights) reaches 55.6% overall under the same protocol.

Usage

The model class lives in the GitHub repository, so clone it first:

git clone https://github.com/mutabalShawandar/tlm-code-refactoring
cd tlm-code-refactoring
pip install -e . huggingface_hub safetensors

Then run a rename edit, which turns cryptic names into meaningful ones:

import json
import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

from src.model_config import ModelConfig
from src.model_variants import TLM_S
from src.tokenizer import TLM_Tokenizer
from src.data_set import apply_sr
from src.inference.generate import decode_rename_without_collision

REPO = "omartra999/tlm-xs-refactor"

# load config, weights and tokenizer from the Hub
cfg = json.load(open(hf_hub_download(REPO, "config.json")))
config = ModelConfig(d_model=cfg["d_model"], d_ff=cfg["d_ff"],
                     n_encoder_layers=cfg["n_encoder_layers"],
                     n_decoder_layers=cfg["n_decoder_layers"])
model = TLM_S(config)
state = load_file(hf_hub_download(REPO, "model.safetensors"))
state["output_head.weight"] = state["embedding.token_embedding.weight"]  # tied weights
model.load_state_dict(state)
model.eval()

tokenizer = TLM_Tokenizer(vocab_size=config.vocab_size)
tokenizer.load(hf_hub_download(REPO, "tokenizer.json"))

code = '''def chk(s):
    if len(s) < 8:
        return False
    return True
'''

# encoder input: <bos> <task token> code <eos>
ids = config.special_token_ids
before_ids = tokenizer.encode(code)
enc_ids = [ids["bos"], ids["rename"]] + before_ids + [ids["eos"]]

with torch.no_grad():
    out = decode_rename_without_collision(model, enc_ids, cfg["task_ids"]["rename"], "cpu",
                                          before_ids, tokenizer, config, collect_conf=False)

# the model emits a SEARCH/REPLACE patch; apply it to the original code
after_ids = apply_sr(before_ids, out["tokens"], config)
print(tokenizer.decode(after_ids))

The model emits a SEARCH/REPLACE patch, which apply_sr applies to the original code. For the chk function above TLM-XS returns the following (the larger models name the parameter password):

def is_valid_password(username):
    if len(username) < 8:
        return False
    return True

(The tokenizer decodes without spaces around operators, so the raw output reads len(username)<8.)

Change the task by swapping "rename" for "simplify", "deduplicate", "remove_dead_code" or "extract_method" (and use greedy_decode from src.inference.generate for the non-rename tasks). Decode with up to 512 new tokens.

Limitations

These models are small and trained on a narrow, partly synthetic corpus. deduplicate and extract_method are the weakest tasks. Outputs should be reviewed and tested before use.