Skip to content

feat: built-in default QRF k-point model with HuggingFace resolution#50

Open
junwen94 wants to merge 1 commit into
stfc:mainfrom
junwen94:feat/qrf-kpoints-reland
Open

feat: built-in default QRF k-point model with HuggingFace resolution#50
junwen94 wants to merge 1 commit into
stfc:mainfrom
junwen94:feat/qrf-kpoints-reland

Conversation

@junwen94

@junwen94 junwen94 commented Jul 8, 2026

Copy link
Copy Markdown
Collaborator

Reland of #38 (reverted in #49) for review. Content unchanged from the cleaned, CI-green state; build artifacts stripped and .gitignore restored. @sigilmakes requesting your architectural review before merge — happy to split it up or discuss design if that helps. Part of #31, gates #30.

Reland of stfc#38 after it was reverted in stfc#49. Content is unchanged from the
cleaned, CI-green state of stfc#38 (build artifacts already stripped, .gitignore
rules restored). Reopening for architectural review before merge.

Part of stfc#31; gates stfc#30.
@junwen94 junwen94 requested a review from sigilmakes July 8, 2026 14:58

@sigilmakes sigilmakes left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hold on, GPT wrote a bad review. Let me write it myself lol

@sigilmakes

Copy link
Copy Markdown
Collaborator

Ok a few issues:

  1. The CLI diverges from the python API by downloading the ML models, whereas the API has no capability for this. This violates the design direction of the CLI (and any frontend interface) being a thin wrapper around the API.

  2. The model downloading needs more thought. At the moment it largely works via side effects in the default_pipeline, which isn't great with no proper model version pinning or documentation that this will even happen. The default model is also under a personal namespace, which would be fine as a temporary measure if it were pinned but ideally needs looking into - some kind of organisation on huggingface for centralising models would be better.

  3. Metallicity seems to be buried inside k-point features and there's still a decent amount of hardcoded behaviour for the QRT advisor. That's fine for now but as we carry down the line we should think about proper unified design for advisors.

  4. CLI tests now run a full pipeline with a real ML model. We absolutely should have tests for that, but for CI/CD we should probably default to a lighter set of tests for day-to-day work and use ML accuracy tests for larger stuff like PRs - though this also depends on how strong the github runners are as well. I doubt they're strong for complex inference workloads.

I'll get my agent to throw in a more detailed analysis here so it can be used as a handoff comment.
Anyway broadly speaking this is good! Just needs some tidying up of the design of systems and the backend stack.

@sigilmakes

sigilmakes commented Jul 8, 2026

Copy link
Copy Markdown
Collaborator

(N.B. This is just a handoff report for agents! Point your agent at this)

Detailed agent architecture handoff for PR #50

I reviewed PR #50 as an architecture/code-health handoff for the QRF reland. Local verification on the reland content passed:

  • uv run pytest → 99 passed
  • uv run ruff check src tests → passed

This branch is directionally useful: the QRF port is readable, the fallback path is mostly honest, and the Kmesh seam is a workable short-term place for structure-dependent k-density prediction. The issues are in the defaulting/model-resolution architecture and in places where QRF-specific behaviour is starting to harden into general Core behaviour.

1. CLI/API default divergence

The CLI now has materially different default behaviour from the Python API.

In src/goldilocks_core/cli/core.py, _pipeline_from_args() returns default_pipeline() when the user supplies no --model, --k-grid, or --k-spacing. That means a bare CLI invocation attempts to use the built-in QRF backend.

But the shared API path remains heuristic by default:

  • run_core_job(..., pipeline=None) constructs Pipeline().
  • goldilocks_core.recommend(...) also calls run_core_job(..., pipeline=None) unless the caller passes a custom pipeline.
  • Pipeline() uses resolve_kpoints_from_advice for the Kmesh stage, not the QRF default.

So:

Entry point Default k-point backend
goldilocks-core recommend file.cif QRF default pipeline
goldilocks_core.recommend(structure) heuristic pipeline
run_core_job(request) heuristic pipeline
Pipeline() heuristic pipeline
default_pipeline() QRF default pipeline

That makes the CLI policy-bearing rather than a thin wrapper around the shared Core API. If the intended product behaviour is “default ML recommendations”, that policy should live in the shared API/config layer, not only in CLI argument handling. If the intended API default remains heuristic, the CLI should not silently choose a different backend.

2. Model downloading/loading is currently a construction-time side effect

src/goldilocks_core/jobs.py::default_pipeline() calls default_kmesh_advisor(), which calls qrf_kdistance_advisor().

qrf_kdistance_advisor() immediately attempts to:

  1. load the QRF model via load_model(spec);
  2. resolve/download metallicity checkpoint and atom-init artifacts;
  3. load the CGCNN metallicity checkpoint.

This happens when the pipeline is constructed, before any structure is loaded and before the Kmesh stage knows whether the request actually needs ML inference.

That has several consequences:

  • constructing a pipeline can perform network/cache/model-loading work;
  • explicit k-point hints bypass inference, but not necessarily advisor construction;
  • CLI unit tests that mock run_core_job can still instantiate a real default model pipeline;
  • local environment state, HF cache state, and pickle/sklearn compatibility can leak into tests that should be pure parse/request tests.

The default pipeline should be side-effect-light. A better shape is a lazy advisor: construct a cheap closure, then resolve/load artifacts only on the first no-hint call that actually requires model inference. Explicit k_grid / k_spacing should bypass both inference and heavyweight model loading.

3. Artifact reproducibility and ownership need tightening

The PR body describes the default as pinned/reproducible, but the code does not pin the Hugging Face revisions.

DEFAULT_KPOINTS_MODEL has:

revision=None

The metallicity artifacts also default to revision=None. That resolves the current default branch/revision on HF, not an immutable model artifact.

There is also an ownership asymmetry:

  • QRF model: STFC-SCD/kpoints-goldilocks-QRF
  • metallicity artifacts: JunwenYin/metallicity-goldilocks-CGCNN

A personal namespace is workable as a temporary development path, but it is not a good long-term default dependency for Core. If the default pipeline depends on this artifact, it should live under an organization-controlled model namespace and be pinned to an immutable revision.

This also matters because ml/models.py::_load_huggingface() downloads the artifact and calls joblib.load(). That makes the trust/reproducibility boundary important. A default path that loads remote pickle/joblib artifacts needs pinned revisions and clear documentation.

4. Download/cache behaviour is undocumented in the CLI UX

A bare CLI invocation can now resolve/download model artifacts, but the help text does not make that clear and there is no obvious explicit escape hatch like --heuristic / --no-model.

Using --k-spacing or --k-grid can avoid the default model path, but that is not the same as a documented backend choice. Users should be able to choose:

  • built-in default ML backend;
  • local model backend;
  • heuristic/no-model backend.

That choice should be explicit in both CLI UX and shared API/config design.

5. QRF feature extraction is hardcoded despite ModelSpec.feature_set

DEFAULT_KPOINTS_MODEL.feature_set records:

qrf_comp_struct_soap_lattice_metal

But qrf_kdistance_advisor() directly imports and calls:

extract_qrf_features(structure, metal_model, atom_init)

So feature_set is metadata only; it is not used for dispatch. That is acceptable for a QRF-specific advisor, but it should not be treated as a general model-backend architecture yet.

As more models land, Core will need a clearer mechanism for mapping model specs to feature extractors and artifact requirements. Otherwise each advisor will hardcode its own model-loading and feature stack.

6. Metallicity is buried inside k-point features

The QRF feature vector requires the CGCNN metallicity representation, so the branch correctly ports enough metallicity machinery to reproduce the trained feature contract.

The architectural risk is that metallicity becomes a hidden implementation detail of k-point prediction. Metallicity is also a structure fact that future Analyze-stage logic may want to expose with provenance/confidence. If that happens later, Core should not end up with two separate metallicity model-loading paths: one buried in QRF features and another in Analyze.

Short-term, this can stay as QRF feature machinery. Long-term, it needs a shared boundary:

  • one artifact resolver/model loader for the metallicity CGCNN;
  • one QRF-specific function that extracts the 64-dim representation for the QRF feature vector;
  • a separate Analyze-facing prediction path if/when metallicity becomes a first-class structure fact.

7. Test strategy needs separation between lightweight CI and real ML validation

The committed tests mostly mock model behaviour, which is good for normal unit tests. But because the CLI can instantiate default_pipeline(), tests can still touch real model-loading paths unintentionally.

There should be a clean split:

  • lightweight unit tests for CLI parsing, request construction, fallback behaviour, and feature assembly with mocks;
  • integration tests that intentionally exercise real HF artifacts / cached model loading / full inference;
  • possibly slower model-quality or accuracy tests that run less frequently, depending on runner capacity.

For day-to-day CI, tests should not accidentally depend on network, HF cache, or heavyweight inference. For PR or scheduled validation, real-model tests are valuable but should be explicit.

8. Smaller code-health issues

These are not the main blockers, but should be cleaned up:

  • contracts.py still says ModelSource supports only local models even though HF support was added.
  • cgcnn.py / metallicity.py describe torch/PyG as optional while this branch declares them as base dependencies.
  • The QRF correction comment/docstring says the correction “widens” the interval. For the default RF95 correction (-0.0016) and formula lower - correction, upper + correction, it narrows the interval. The implementation matches the reference formula; the wording/test coverage should be fixed.
  • Add a direct test for full extract_qrf_features() assembly: 419 structure features plus a mocked 64-dim metallicity block equals 483 features, with matching feature names.
  • Consider validating QRF outputs before converting the median to a mesh: finite values, positive median, and sane interval ordering.

Summary

The implementation is not sloppy: the QRF path is readable, local checks pass, fallback provenance is mostly good, and the existing Kmesh stage can host structure-dependent k-density prediction for now.

The reland should still change before merge because it currently makes default backend policy, remote artifact resolution, model pinning, and CLI/API layering decisions implicitly. The important fixes are:

  1. unify CLI and Python API default policy;
  2. make default pipeline/advisor construction lazy and side-effect-light;
  3. pin HF artifacts and move core defaults to organization-controlled namespaces;
  4. document model download/cache behaviour and expose an explicit heuristic/no-model backend choice;
  5. keep the QRF hardcoding clearly scoped while planning the unified advisor/model-feature design.

Written by an agent on behalf of Willow.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants