File size: 4,540 Bytes
43a90a1 | 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 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | import json
import numpy as np
import pytest
@pytest.fixture(scope="module")
def tok():
from needle.model.tokenizer import get_tokenizer
return get_tokenizer(8192)
def test_render_example_markers():
from needle.model.finetune import render_example
from needle.model.tokenizer import IM_START, TOOLS_START, TOOL_CALL_START, THINK_START
example = {"tools": [{"name": "f", "parameters": {"type": "object", "properties": {}}}],
"query": "do the thing", "reasoning": "because reasons",
"answers": [{"name": "f", "arguments": {}}]}
prompt, target = render_example(example)
assert IM_START in prompt
assert TOOLS_START in prompt
assert "do the thing" in prompt
assert THINK_START in target
assert "because reasons" in target
assert TOOL_CALL_START in target
assert '"name":"f"' in target
def test_render_example_matches_training_template():
from needle.model.finetune import render_example
from needle.model.tokenizer import (IM_END, THINK_START, THINK_END,
TOOL_CALL_START, TOOL_CALL_END)
example = {"tools": [], "query": "q", "reasoning": "'q' -> f",
"answers": [{"name": "f", "arguments": {}}]}
_, target = render_example(example)
assert target == (THINK_START + "\n'q' -> f\n" + THINK_END + "\n"
+ TOOL_CALL_START + '[{"name":"f","arguments":{}}]'
+ TOOL_CALL_END + IM_END)
def test_render_example_empty_reasoning_omits_think():
from needle.model.finetune import render_example
from needle.model.tokenizer import THINK_START, TOOL_CALL_START
for example in ({"tools": [], "query": "q", "answers": []},
{"tools": [], "query": "q", "reasoning": "", "answers": []},
{"tools": [], "query": "q", "reasoning": None, "answers": []},
{"tools": [], "query": "q", "reasoning": " ", "answers": []}):
_, target = render_example(example)
assert THINK_START not in target
assert target.startswith(TOOL_CALL_START)
def test_render_example_keeps_unicode():
from needle.model.finetune import render_example
example = {"tools": [{"name": "f", "parameters": {"type": "object", "properties": {
"name": {"type": "string", "description": "café"}}}}],
"query": "add Tomás",
"answers": [{"name": "f", "arguments": {"name": "Tomás"}}]}
prompt, target = render_example(example)
assert "café" in prompt
assert '"Tomás"' in target
def test_render_example_accepts_function_calls_alias():
from needle.model.finetune import render_example
example = {"tools": [], "query": "hi",
"function_calls": [{"name": "g", "arguments": {"x": 1}}]}
_, target = render_example(example)
assert '"name":"g"' in target
def test_encode_loss_mask_targets_only(tok):
from needle.model.finetune import _encode, render_example
from needle.model.tokenizer import BOS_ID
example = {"tools": [{"name": "f", "parameters": {"type": "object", "properties": {}}}],
"query": "hello", "answers": [{"name": "f", "arguments": {}}]}
ids, mask = _encode(tok, example, max_len=128)
assert len(ids) == 128 and len(mask) == 128
assert ids[0] == BOS_ID
assert mask[0] == 0.0
assert any(m == 1.0 for m in mask)
_, target = render_example(example)
assert sum(1 for m in mask if m == 1.0) == len(tok.encode(target)) + 1
def test_fit_max_len_buckets(tok, tmp_path):
from needle.model.finetune import fit_max_len
path = tmp_path / "data.jsonl"
with open(path, "w") as handle:
handle.write(json.dumps({"tools": [], "query": "short", "answers": []}) + "\n")
assert fit_max_len(str(path), tok, 1024) == 128
assert fit_max_len(str(path), tok, 64) == 64
def test_load_jsonl_shapes_and_skips_invalid(tok, tmp_path):
from needle.model.finetune import load_jsonl
path = tmp_path / "data.jsonl"
rows = [
{"tools": [], "query": "a", "answers": []},
{"tools": [], "query": "b", "answers": []},
{"tools": [], "reasoning": "no query here"},
]
with open(path, "w") as handle:
for row in rows:
handle.write(json.dumps(row) + "\n")
handle.write("\n")
seqs, masks = load_jsonl(str(path), tok, max_len=32)
assert seqs.shape == (2, 32)
assert masks.shape == (2, 32)
assert seqs.dtype == np.int32
assert masks.dtype == np.float32
|