Spaces:
Sleeping
Sleeping
File size: 2,047 Bytes
2d52135 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | from pathlib import Path
from fastapi.testclient import TestClient
from mini_transformer.apps import server
from mini_transformer.model_loader import MODEL_NAME_ENV, MODELS_ENV
MINIMAL_CFG = """
model:
name: demo
best_checkpoint_path: ./checkpoints/best.pt
latest_checkpoint_path: null
tokenizer: demo_tok
d_model: 16
num_layers: 1
num_heads: 2
d_ff: 32
dropout_rate: 0.0
vocab_size: 32
max_seq_len: 32
pad_id: 0
bos_id: 1
eos_id: 2
tokenizer:
name: demo_tok
path: ./tokenizer/tokenizer.json
corpus: null
dataset: null
vocab_size: 32
max_seq_len: 32
pad_token: "<pad>"
bos_token: "<bos>"
eos_token: "<eos>"
unk_token: "<unk>"
special_tokens: ["<pad>", "<bos>", "<eos>", "<unk>"]
pad_id: 0
bos_id: 1
eos_id: 2
unk_id: 3
generation:
max_new_tokens: 4
temperature: 1.0
top_k: null
top_p: null
do_sample: false
presence_penalty: 0.0
frequency_penalty: 0.0
no_repeat_ngram: null
min_steps_before_eos: 0
runtime:
seed: 42
device: cpu
output_dir: ./outputs
data_dir: ./data
tokenizer_dir: ./tokenizer
cache_dir: ./cache
checkpoint_path: ./checkpoints
input_text: ""
"""
def _make_demo_model(tmp_path: Path) -> Path:
models_root = tmp_path / "trained_models"
model_dir = models_root / "demo"
config_dir = model_dir / "configs"
config_dir.mkdir(parents=True, exist_ok=True)
(model_dir / "checkpoints").mkdir(exist_ok=True)
(model_dir / "tokenizer").mkdir(exist_ok=True)
(config_dir / "config_inference.yaml").write_text(MINIMAL_CFG)
return models_root
def test_generate_endpoint_uses_local_model(monkeypatch, tmp_path):
models_root = _make_demo_model(tmp_path)
monkeypatch.setenv(MODELS_ENV, str(models_root))
monkeypatch.setenv(MODEL_NAME_ENV, "demo")
monkeypatch.setattr(server, "run_inference", lambda cfg: ["ok"])
client = TestClient(server.app)
response = client.get("/generate", params={"text": "hello"})
assert response.status_code == 200
assert response.json()["outputs"] == ["ok"]
|