File size: 1,382 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 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 | import torch
def test_generation_accepts_prompt(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=32,
temperature=0.0,
)
assert generated.shape[1] == len(prompt) + 32
def test_generated_output_is_valid_text(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=32,
temperature=0.0,
)
text = tokenizer.decode(generated[0].tolist())
assert isinstance(text, str)
assert text.startswith(prompt)
assert len(text) == len(prompt) + 32
def test_generation_is_deterministic_with_seed(model, tokenizer):
prompt = "Wherefore art thou"
token_ids = torch.tensor(
tokenizer.encode(prompt),
dtype=torch.long,
).unsqueeze(0)
torch.manual_seed(7)
first = model.generate(
token_ids=token_ids,
max_new_tokens=32,
temperature=0.8,
)
torch.manual_seed(7)
second = model.generate(
token_ids=token_ids,
max_new_tokens=32,
temperature=0.8,
)
assert torch.equal(first, second) |