File size: 561 Bytes
f0b317c | 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 | import torch
def test_model_loads_on_cpu(model):
assert next(model.parameters()).device.type == "cpu"
def test_cpu_generation_completes(model, tokenizer):
prompt = "ROMEO:"
token_ids = torch.tensor(
tokenizer.encode(prompt),
dtype=torch.long,
).unsqueeze(0)
generated = model.generate(
token_ids=token_ids,
max_new_tokens=64,
temperature=0.8,
top_k=40,
)
text = tokenizer.decode(generated[0].tolist())
assert len(text) == len(prompt) + 64
assert text.startswith(prompt) |