mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-17 16:37:30 +02:00
50 lines
1.4 KiB
Python
50 lines
1.4 KiB
Python
"""Pinned, networked tiny-model load plus cache-only replay."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import uuid
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from obliteratus.models.loader import load_model
|
|
|
|
|
|
pytestmark = [pytest.mark.network, pytest.mark.download]
|
|
MODEL = "hf-internal-testing/tiny-random-gpt2"
|
|
REVISION = "71034c5d8bde858ff824298bdedc65515b97d2b9"
|
|
|
|
|
|
def test_pinned_tiny_model_download_inference_and_offline_cache(monkeypatch):
|
|
handle = load_model(
|
|
MODEL,
|
|
revision=REVISION,
|
|
device="cpu",
|
|
dtype="float32",
|
|
trust_remote_code=False,
|
|
skip_snapshot=True,
|
|
)
|
|
encoded = handle.tokenizer("conditional gate", return_tensors="pt")
|
|
with torch.no_grad():
|
|
output = handle.model(**encoded)
|
|
assert output.logits.shape[:2] == encoded["input_ids"].shape
|
|
assert next(handle.model.parameters()).device.type == "cpu"
|
|
handle.cleanup()
|
|
|
|
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
|
|
monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1")
|
|
cached = load_model(
|
|
MODEL,
|
|
revision=REVISION,
|
|
device="cpu",
|
|
local_files_only=True,
|
|
trust_remote_code=False,
|
|
skip_snapshot=True,
|
|
)
|
|
assert cached.model_name == MODEL
|
|
cached.cleanup()
|
|
|
|
missing = f"obliteratus/offline-missing-{uuid.uuid4().hex}"
|
|
with pytest.raises(OSError):
|
|
load_model(missing, revision=REVISION, device="cpu", local_files_only=True)
|