Skip to content

Python API

The BasinSurgeon class is the primary Python interface for Energy Language.

Construction

python
from qriton_hlm import BasinSurgeon

# From a checkpoint file
surgeon = BasinSurgeon.from_checkpoint("model.pt", device="cuda")

# From a raw weight matrix
surgeon = BasinSurgeon.from_W(W_tensor, device="cpu")

Basin Operations

python
# Survey all basins in a layer
survey = surgeon.survey(layer=0)
# Returns: {num_basins, basins, energies, populations, ...}

# Inject a new basin
result = surgeon.inject(layer=0, seed=42, strength=0.1)

# Remove a basin
result = surgeon.remove(layer=0, seed=17, strength=0.1)

# Move a basin
result = surgeon.move(layer=0, seed=42, strength=0.1)

# Verify a basin exists
v = surgeon.verify(layer=0, seed=42)
# Returns: {is_basin, cos, energy, iters}

# Strengthen / weaken
surgeon.strengthen(layer=0, seed=42, factor=1.5)
surgeon.weaken(layer=0, seed=42, factor=0.5)

# Measure energy at a point
e = surgeon.energy(layer=0, seed=42)

Concept Operations

python
# Capture concept from text (requires full model)
surgeon.capture(layer=5, text="Thank you so much", concept_name="polite")
surgeon.capture(layer=5, text="I really appreciate it", concept_name="polite")

# Batch capture
surgeon.batch_capture(layer=5, concept_name="polite", texts=[
    "Thank you so much",
    "I really appreciate it",
    "That's very kind of you",
])

# Inject captured concept
result = surgeon.inject_concept(layer=5, concept_name="polite", strength=0.1)

# Remove concept
surgeon.remove_concept(layer=5, concept_name="polite", strength=0.1)

# Blend two concepts
surgeon.blend("polite", "formal", "professional", ratio=0.6)

# List all concepts
surgeon.list_concepts()

# Export / Import
surgeon.export_concept("polite", "polite.concept")
surgeon.import_concept("polite.concept")

# Transplant from another model
other = BasinSurgeon.from_checkpoint("other_model.pt")
surgeon.transplant(other, layer=5, concept_name="humor")

Persistence

python
# Apply modified W to live model
surgeon.apply(layer=5)

# Restore original W
surgeon.restore(layer=5)
surgeon.restore_all()

# Check what changed
diff = surgeon.diff(layer=5)
# Returns: {frobenius_norm, relative_pct, max_change, ...}

# Compare with another model
diff = surgeon.compare(other_surgeon, layer=5)

# View operation history
surgeon._history

Low-Level Primitives

For direct tensor operations:

python
from qriton_hlm import (
    poly_interaction,
    compute_energy,
    find_basins,
    inject_basin,
    remove_basin,
    move_basin,
    verify_basin_exists,
    load_W_from_checkpoint,
)

# Example: manual basin injection
W_new = inject_basin(W, target_vector, strength=0.1, beta=7.0)
exists, state, cos, iters = verify_basin_exists(W_new, target_vector)