Code LLM Training › Implementation: FIM Transformation

Michael BrenndoerferJanuary 29, 20263 min read
Language AI HandbookAll chapters
Part 47 of 67Code Generation
Chapter 357Code LLM Training › Implementation: FIM Transformation
Reading section 8 of 13 · Implementation: FIM Transformation · 3 minSection 8 of 13 · 3 min

Book contents

Language AI Handbook67 parts · 513 chaptersView the whole book →
Active chapter 357 of 513

Section 8 of 13 · 3 min read

Code LLM Training › Implementation: FIM TransformationOne section · Sections in this chapter8 / 13 ↓
Sections in this chapter13
Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Now let's implement the fill-in-the-middle transformation. We'll apply it to a realistic code snippet and show the resulting training format for both the PSM and SPM variants.

In[9]:
Code
# FIM special tokens
PRE_TOKEN = "<PRE>"
SUF_TOKEN = "<SUF>"
MID_TOKEN = "<MID>"
EOT_TOKEN = "<EOT>"  # end of text


def fim_transform_psm(document: str, split_rate: float = 0.5) -> str:
    """
    Apply PSM (Prefix-Suffix-Middle) FIM transformation.

    With probability (1 - split_rate), return the document unchanged (causal LM format).
    With probability split_rate, apply FIM transformation.

    Returns the transformed document as a single string.
    """
    if random.random() > split_rate:
        return document  # Standard causal LM training

    lines = document.split("\n")
    if len(lines) < 3:
        return document  # Too short to split meaningfully

    # Choose random line boundaries for prefix/suffix split
    split_line = random.randint(1, len(lines) - 2)

    prefix = "\n".join(lines[:split_line])
    middle = lines[split_line]
    suffix = "\n".join(lines[split_line + 1 :])

    # PSM format: PRE prefix SUF suffix MID middle EOT
    return (
        f"{PRE_TOKEN}{prefix}{SUF_TOKEN}{suffix}{MID_TOKEN}{middle}{EOT_TOKEN}"
    )


def fim_transform_spm(document: str, split_rate: float = 0.5) -> str:
    """
    Apply SPM (Suffix-Prefix-Middle) FIM transformation.

    SPM puts the suffix before the prefix, which some research shows
    improves infilling quality by conditioning generation on the suffix first.
    """
    if random.random() > split_rate:
        return document

    lines = document.split("\n")
    if len(lines) < 3:
        return document

    split_line = random.randint(1, len(lines) - 2)

    prefix = "\n".join(lines[:split_line])
    middle = lines[split_line]
    suffix = "\n".join(lines[split_line + 1 :])

    # SPM format: PRE SUF suffix MID prefix middle EOT
    return f"{PRE_TOKEN}{SUF_TOKEN}{suffix}{MID_TOKEN}{prefix}\n{middle}{EOT_TOKEN}"
In[10]:
Code
# Apply FIM to a specific example to show the transformation
example_code = deduped_corpus[0]

# Force transformation for demonstration (override random)
random.seed(0)
transformed_psm = fim_transform_psm(example_code, split_rate=1.0)

random.seed(0)
transformed_spm = fim_transform_spm(example_code, split_rate=1.0)
Out[11]:
Console
=== ORIGINAL CODE ===
import numpy as np

def compute_mean(values):
    """Compute the arithmetic mean of a list."""
    return np.mean(values)

result = compute_mean([1, 2, 3, 4, 5])
print(result)

=== PSM TRANSFORMED ===
<PRE>
import numpy as np

def compute_mean(values):
    """Compute the arithmetic mean of a list."""
    return np.mean(values)
<SUF>print(result)
<MID>result = compute_mean([1, 2, 3, 4, 5])<EOT>
...

=== SPM TRANSFORMED ===
<PRE><SUF>print(result)
<MID>
import numpy as np

def compute_mean(values):
    """Compute the arithmetic mean of a list."""
    return np.mean(values)

result = compute_mean([1, 2, 3, 4, 5])<EOT>
...

The PSM and SPM formats are both valid training examples. The model learns to recognize the <PRE>, <SUF>, and <MID> sentinel tokens as instructions for how to interpret the surrounding context.

In[12]:
Code
# Build the final training batch with mixed FIM and causal LM examples
def build_training_batch(corpus: list[str], fim_rate: float = 0.5) -> list[str]:
    """
    Transform each document with probability fim_rate using FIM.
    The remaining documents are kept in standard causal LM order.
    """
    batch = []
    for doc in corpus:
        if random.random() < fim_rate:
            # Randomly choose PSM or SPM
            if random.random() < 0.5:
                batch.append(fim_transform_psm(doc, split_rate=1.0))
            else:
                batch.append(fim_transform_spm(doc, split_rate=1.0))
        else:
            batch.append(doc)  # Standard causal LM
    return batch


random.seed(42)
training_batch = build_training_batch(deduped_corpus, fim_rate=0.5)
Out[13]:
Console
Training batch composition:
  FIM examples:        5 (71%)
  Causal LM examples:  2 (29%)
  Total:               7

The training batch is now a mix of FIM and standard causal LM examples. The model learns to handle both: standard left-to-right completion from prefix-only contexts, and infilling when prefix and suffix are both present.