We are going to build an LLM training assistant that generates PyTorch implementations of transformer components and explains the math behind them. If you are learning how LLMs work under the hood, this gives you a reusable coding partner that runs on Oxlo.ai's flat-rate API. The finished script is under 80 lines and uses only the OpenAI SDK.
What you'll need
- Python 3.10 or newer
pip install openai- An Oxlo.ai API key from https://portal.oxlo.ai
Step 1: Make a first call to a coding model
Before adding abstractions, verify that Oxlo.ai returns valid PyTorch with a single user message. I use deepseek-v3.2 because it handles reasoning and code generation well.
from openai import OpenAI
client = OpenAI(base_url="https://api.oxlo.ai/v1", api_key="YOUR_OXLO_API_KEY")
response = client.chat.completions.create(
model="deepseek-v3.2",
messages=[
{"role": "user", "content": "Write a minimal PyTorch implementation of multi-head self-attention with einsum. No main function."},
],
)
print(response.choices[0].message.content)Step 2: Lock in a system prompt
A generic chat model drifts between teaching styles. We pin it to a senior ML engineer persona who returns only code and a short explanation. Edit the prompt below to change the agent's personality or constraints.
from openai import OpenAI
client = OpenAI(base_url="https://api.oxlo.ai/v1", api_key="YOUR_OXLO_API_KEY")
SYSTEM_PROMPT = """You are a senior ML engineer helping a colleague build an LLM from scratch in PyTorch.
Rules:
- Provide only valid, runnable PyTorch code.
- After the code, give a 2-sentence explanation of the key mathematical idea.
- Do not use external libraries beyond torch and torch.nn.
- If the user asks for a variant, rewrite the full function so it remains self-contained."""
response = client.chat.completions.create(
model="deepseek-v3.2",
messages=[
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": "Implement scaled dot-product attention with causal masking."},
],
)
print(response.choices[0].message.content)Step 3: Enforce structured output with JSON mode
Parsing free-form markdown is brittle. Oxlo.ai supports JSON mode, so we request a strict schema with 'code' and 'explanation' fields that our script can consume directly.
from openai import OpenAI
import json
client = OpenAI(base_url="https://api.oxlo.ai/v1", api_key="YOUR_OXLO_API_KEY")
SYSTEM_PROMPT = """You are a senior ML engineer helping a colleague build an LLM from scratch in PyTorch.
Rules:
- Provide only valid, runnable PyTorch code.
- After the code, give a 2-sentence explanation of the key mathematical idea.
- Do not use external libraries beyond torch and torch.nn.
- If the user asks for a variant, rewrite the full function so it remains self-contained."""
response = client.chat.completions.create(
model="deepseek-v3.2",
messages=[
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": "Implement scaled dot-product attention with causal masking. Return JSON with keys 'code' and 'explanation'."},
],
response_format={"type": "json_object"},
)
result = json.loads(response.choices[0].message.content)
print("=== CODE ===")
print(result["code"])
print("=== EXPLANATION ===")
print(result["explanation"])Step 4: Iterate with multi-turn context
Real architectures are built in layers. We keep conversation history so the model can revise the previous block, such as adding Rotary Position Embedding to the attention mechanism it just wrote.
from openai import OpenAI
import json
client = OpenAI(base_url="https://api.oxlo.ai/v1", api_key="YOUR_OXLO_API_KEY")
SYSTEM_PROMPT = """You are a senior ML engineer helping a colleague build an LLM from scratch in PyTorch.
Rules:
- Provide only valid, runnable PyTorch code.
- After the code, give a 2-sentence explanation of the key mathematical idea.
- Do not use external libraries beyond torch and torch.nn.
- If the user asks for a variant, rewrite the full function so it remains self-contained."""
messages = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": "Implement scaled dot-product attention with causal masking. Return JSON with keys 'code' and 'explanation'."},
]
response = client.chat.completions.create(
model="deepseek-v3.2",
messages=messages,
response_format={"type": "json_object"},
)
result = json.loads(response.choices[0].message.content)
messages.append({"role": "assistant", "content": json.dumps(result)})
messages.append({"role": "user", "content": "Now rewrite that same function to use Rotary Position Embedding (RoPE) instead of absolute positional encodings. Return JSON."})
response = client.chat.completions.create(
model="deepseek-v3.2",
messages=messages,
response_format={"type": "json_object"},
)
result = json.loads(response.choices[0].message.content)
print(result["code"])Step 5: Wrap everything in a reusable helper
Copy-pasting client setup for every component gets old. This helper accepts a conversation history and returns the parsed result, so you can assemble an entire transformer block by block.
from openai import OpenAI
import json
client = OpenAI(base_url="https://api.oxlo.ai/v1", api_key="YOUR_OXLO_API_KEY")
SYSTEM_PROMPT = """You are a senior ML engineer helping a colleague build an LLM from scratch in PyTorch.
Rules:
- Provide only valid, runnable PyTorch code.
- After the code, give a 2-sentence explanation of the key mathematical idea.
- Do not use external libraries beyond torch and torch.nn.
- If the user asks for a variant, rewrite the full function so it remains self-contained."""
def generate_component(messages, model="deepseek-v3.2"):
response = client.chat.completions.create(
model=model,
messages=messages,
response_format={"type": "json_object"},
)
return json.loads(response.choices[0].message.content)
history = [
{"role": "system", "content": SYSTEM_PROMPT},
]
prompts = [
"Write a PyTorch RMSNorm layer. Return JSON with keys 'code' and 'explanation'.",
"Write a PyTorch SwiGLU feed-forward layer compatible with Llama-style dims. Return JSON.",
"Write a PyTorch DecoderBlock that composes RMSNorm, causal self-attention, and SwiGLU. Return JSON.",
]
for p in prompts:
history.append({"role": "user", "content": p})
out = generate_component(history)
history.append({"role": "assistant", "content": json.dumps(out)})
print(f"--- {p[:40]}... ---")
print(out["code"][:300] + "...")Run it
Here is a complete end-to-end script that generates a minimal GPT-style architecture. I run it against Oxlo.ai and show the first few lines of the generated model.
from openai import OpenAI
import json
client = OpenAI(base_url="https://api.oxlo.ai/v1", api_key="YOUR_OXLO_API_KEY")
SYSTEM_PROMPT = """You are a senior ML engineer helping a colleague build an LLM from scratch in PyTorch.
Rules:
- Provide only valid, runnable PyTorch code.
- After the code, give a 2-sentence explanation of the key mathematical idea.
- Do not use external libraries beyond torch and torch.nn.
- If the user asks for a variant, rewrite the full function so it remains self-contained."""
def generate_component(messages, model="qwen-3-32b"):
response = client.chat.completions.create(
model=model,
messages=messages,
response_format={"type": "json_object"},
)
return json.loads(response.choices[0].message.content)
history = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": "Generate a complete, minimal GPT-style LM in PyTorch with an Embedding layer, 2 DecoderBlocks, and a head that outputs logits over a vocab of 1000. Return JSON with keys 'code' and 'explanation'."},
]
result = generate_component(history)
print(result["code"])Example output (truncated):
class GPT(nn.Module):
def __init__(self, vocab_size=1000, dim=512, heads=8, layers=2):
super().__init__()
self.tok_emb = nn.Embedding(vocab_size, dim)
self.blocks = nn.ModuleList([DecoderBlock(dim, heads) for _ in range(layers)])
self.norm = RMSNorm(dim)
self.head = nn.Linear(dim, vocab_size, bias=False)
def forward(self, x):
...
Explanation: This stacks two decoder blocks with causal self-attention and SwiGLU, followed by a linear projection to vocabulary logits, matching the minimal GPT architecture.
Wrap-up
You now have a scriptable assistant that generates transformer components on demand. Two concrete next steps: add a local torch.jit.script or exec sanity check so the agent's code is validated before you save it to disk, and swap in llama-3.3-70b or kimi-k2.6 when you move from block generation to full training-loop design. Because Oxlo.ai charges a flat rate per request, you can send long system prompts and multi-turn histories without watching token meters climb.
Top comments (0)