TrainLens

Turn AI training runs into research-grade notebook reports.

TrainLens is a Jupyter-first support package for AI model training and research workflows. Its main value is that you can turn the training state already present in your notebook into structured analysis without leaving Jupyter, opening a separate chat, or copying experiment context by hand.

It redacts likely secrets and asks an OpenAI-compatible LLM to draft:

• a scientific paper-style training report
• an evidence-backed improvement plan

It is useful when you run many experiments and want consistent explanations of metrics, datasets, hyperparameters, limitations, and next steps directly inside the notebook where the work is happening.

Actively maintained by Eduardo J. Barrios.
For issues or improvements, open a pull request or email edujbarrios@outlook.com.

Quickstart

Install TrainLens:

pip install trainlens

Configure an OpenAI-compatible LLM provider:

import os

os.environ["TRAINLENS_LLM_BASE_URL"] = "https://api.openai.com/v1"
os.environ["TRAINLENS_LLM_API_KEY"] = "your-api-key"
os.environ["TRAINLENS_LLM_MODEL"] = "gpt-5.4-mini"

Keep useful experiment state in ordinary notebook variables:

dataset_name = "synthetic_cpu_binary_classification"
dataset_notes = "240 synthetic 2D samples; balanced validation split."
model_name = "pure-python-logistic-regression"
training_params = {"epochs": 30, "learning_rate": 0.45}

history = {
    "train_loss": [0.6095, 0.5005, 0.2459],
    "eval_loss": [0.6164, 0.5160, 0.2808],
    "accuracy": [0.8889, 0.9056, 0.9222],
    "val_accuracy": [0.8167, 0.8333, 0.9167],
}

Run the notebook magics:

%load_ext trainlens.magic.extension

# Scientific paper-style training report
%explain_training

# Evidence-backed experiment plan
%suggest_improvements

Local Models

TrainLens can avoid external API token costs by using a local OpenAI-compatible server. Example with Ollama:

ollama pull llama3.1
ollama serve

import os

os.environ["TRAINLENS_LLM_BASE_URL"] = "http://localhost:11434/v1"
os.environ["TRAINLENS_LLM_API_KEY"] = "ollama"
os.environ["TRAINLENS_LLM_MODEL"] = "llama3.1"

The same pattern works with LM Studio, vLLM, llama.cpp server, and similar local servers.

Python API

from trainlens import build_improvement_ideas, build_paper_report

paper = build_paper_report()
ideas = build_improvement_ideas()

Inside Jupyter, helpers read the active notebook namespace automatically. Outside IPython, pass an explicit dictionary-like namespace.

Token Usage

Costs depend on provider, model, tokenizer, prompt size, and output length. These are rough estimates, not billing records.

Approximate USD cost for a 50% input / 50% output token mix, using public prices checked on 2026-07-07:

Cost formula:

\text{estimated cost} =
\frac{
\text{input tokens} \times \text{input price per 1M}
+ \text{output tokens} \times \text{output price per 1M}
}{1{,}000{,}000}

Privacy

TrainLens summarizes visible notebook state, redacts likely secrets, truncates large literals, and sends only the resulting context to the configured endpoint. Do not store real API keys in notebooks or committed files.

License

Apache License 2.0. See LICENSE.