pycsamt.agents.model_zoo_agent#

pycsamt.agents.model_zoo_agent#

ModelZooAgent — Browse, download, and deploy pre-trained EM inverters.

The model zoo is hosted at github.com/earthai-tech/pycsamt-models and managed by pycsamt.ai._zoo. Checkpoints are cached locally in ~/.pycsamt/model_zoo/ (overridden via PYCSAMT_MODEL_CACHE).

Three operations#

"list" (default)

Print the full registry — model names, architectures, descriptions. No network access required.

"download"

Download a named checkpoint to the local cache and return its path.

"predict"

Download → load into the matching EMInverter1D → predict on observed sites and plot the resistivity section. This is the Phase 5 model zoo shortcut — users get a production-quality inverter in one call, without training from scratch.

Pre-trained model naming convention#

<method>-<arch>-<n_layers>layer-v<version>

Examples: mt1d-resnet-5layer-v1, mt1d-cnn-5layer-v1, csamt1d-resnet-5layer-v1, tem1d-fcn-5layer-v1.

Notes

  • Weights are released in Phase 5. Until then, download_checkpoint() raises a RuntimeError explaining the situation.

  • The "predict" action falls back gracefully to on-the-fly training when the checkpoint is unavailable.

Classes

ModelZooAgent(*[, api_key, model, ...])

Browse, download, and run pre-trained EM inverters from the model zoo.
class pycsamt.agents.model_zoo_agent.ModelZooAgent(*, api_key=None, model=None, llm_provider='claude', cache_dir=None, force_download=False)[source]#

Bases: BaseAgent

Browse, download, and run pre-trained EM inverters from the model zoo.

Parameters:

  • api_key (str)

  • model (model_info dict — zoo metadata for the requested)

  • llm_provider (str)

  • cache_dir (str or None) – Override default cache ~/.pycsamt/model_zoo/.

  • force_download (bool) – Re-download even if cached (default False).

  • keys (Output data)

  • ----------

  • action (str) – "list" (default), "download", or "predict".

  • model_name (str) – Required for "download" and "predict" actions. E.g. "mt1d-resnet-5layer-v1".

  • path (sites /) – Required for "predict" action.

  • output_dir (str, optional)

  • keys

  • ----------------

  • performed (action str — which action was)

  • (action="list") (models dict — full registry)

  • (action="download"/"predict") (``checkpoint_path``str — local path) –

  • model

  • {station (predictions dict —)

  • (action="predict") (rms_global float)

  • dict (figure_paths)

  • dict

Examples

List available models:

agent = ModelZooAgent()
r = agent.execute({"action": "list"})
for name, desc in r["models"].items():
    print(name, "—", desc[:60])

Download a checkpoint:

r = agent.execute({"action": "download", "model_name": "mt1d-resnet-5layer-v1"})
print(r["checkpoint_path"])

Predict on observed sites (fine-tune skipped if weights unavailable):

r = agent.execute({
    "action":     "predict",
    "model_name": "mt1d-resnet-5layer-v1",
    "path":       "/data/WILLY_EDIs",
    "output_dir": "/out/zoo_predict",
})
print(r["rms_global"])
SYSTEM_PROMPT: str = 'You are an expert in pre-trained geophysical AI models.\nGiven a model zoo query result, write 2-3 sentences that:\n1. Describe which pre-trained model was used and its provenance (architecture, training data).\n2. Comment on the prediction quality (RMS, reliability) relative to the expected use case.\n3. Recommend whether the user should fine-tune on their own data or use the pre-trained weights directly.\nReply in plain English.\n'#

Override in subclasses to give the LLM its domain expertise.

execute(input_data)[source]#

Run this agent on input_data and return an AgentResult.

Subclasses must implement this method. The contract:

  • Reset self._last_cost = 0.0 at the top.

  • Record wall-clock time with t0 = time.time().

  • Return AgentResult(elapsed_seconds=time.time()-t0, cost_estimate_usd=self._last_cost, ...).

Parameters:

input_data (dict[str, Any])

Return type:

AgentResult