| import hashlib |
| import importlib.util |
| import struct |
| import sys |
| from pathlib import Path |
|
|
| RELEASE_MODULE_NAME = "clef_release_joint_schema_model" |
| RELEASE_FILE = "joint_schema_model.py" |
| _release = None |
|
|
|
|
| def release_module(snapshot=None): |
| global _release |
| if _release is not None and (snapshot is None or _release.__file__ == str(Path(snapshot) / RELEASE_FILE)): |
| return _release |
| if snapshot is None: |
| from models.autoports.cloudflare_clef.tt.loader import resolve_snapshot |
|
|
| snapshot = resolve_snapshot() |
| path = Path(snapshot) / RELEASE_FILE |
| spec = importlib.util.spec_from_file_location(RELEASE_MODULE_NAME, str(path)) |
| module = importlib.util.module_from_spec(spec) |
| sys.modules[RELEASE_MODULE_NAME] = module |
| spec.loader.exec_module(module) |
| _release = module |
| return module |
|
|
|
|
| def encode(tokenizer, record, processor=None, max_length=16384, max_state_tokens=None): |
| return release_module().encode_record( |
| tokenizer, record, max_length=max_length, max_state_tokens=max_state_tokens, processor=processor |
| ) |
|
|
|
|
| def schema_ids_and_question_starts(tokenizer, record): |
| m = release_module() |
| tokens = m._tokens |
| schema_ids = tokens(tokenizer, "\n\nSCHEMA FIELDS:\n") |
| question_starts = [] |
| for question_index, (question_id, question) in enumerate(record["questions"].items()): |
| schema_ids.extend( |
| tokens( |
| tokenizer, |
| f"\nFIELD {question_index + 1}\nID: {question_id}\nTYPE: {question['type']}\nINSTRUCTION: ", |
| ) |
| ) |
| question_starts.append(len(schema_ids)) |
| instructions = question.get("instructions") |
| if instructions is None or instructions == "": |
| instructions = str(question_id) |
| schema_ids.extend(tokens(tokenizer, m.render(instructions))) |
| schema_ids.extend(tokens(tokenizer, "\nALLOWED OPTIONS:\n")) |
| for option_index, (option_id, description) in enumerate(m.question_options(question)): |
| schema_ids.extend(tokens(tokenizer, f"OPTION {option_index + 1}: ")) |
| semantics = {"option_id": option_id} |
| if description is not None: |
| semantics["description"] = description |
| schema_ids.extend(tokens(tokenizer, m.render(semantics))) |
| schema_ids.extend(tokens(tokenizer, "\n")) |
| schema_ids.extend(tokens(tokenizer, "END FIELD\n")) |
| return schema_ids, question_starts |
|
|
|
|
| def prefix_ids(tokenizer): |
| m = release_module() |
| return m._tokens(tokenizer, f"<|im_start|>system\n{m.SYSTEM_PROMPT}<|im_end|>\n<|im_start|>user\nSTATE:\n") |
|
|
|
|
| def suffix_ids(tokenizer): |
| m = release_module() |
| return m._tokens(tokenizer, "\n<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\nJOINT SCHEMA DECISIONS:") |
|
|
|
|
| def split_for_cache(encoded, tokenizer, record): |
| ids = list(encoded.input_ids) |
| schema, question_starts = schema_ids_and_question_starts(tokenizer, record) |
| suffix = suffix_ids(tokenizer) |
| tail = schema + suffix |
| split = len(ids) - len(tail) |
| if split < 0 or ids[split:] != tail: |
| raise ValueError("schema and suffix ids do not match the tail of the encoded record") |
| head = prefix_ids(tokenizer) |
| if ids[: len(head)] != head: |
| raise ValueError("prefix ids do not match the head of the encoded record") |
| if len(question_starts) != len(encoded.questions): |
| raise ValueError("question count differs between the schema recomputation and the encoded record") |
| for start, question in zip(question_starts, encoded.questions): |
| if question.question_span[0] != split + start: |
| raise ValueError(f"question {question.question_id} span does not start at the recomputed offset") |
| return ids[:split], ids[split:], encoded.questions |
|
|
|
|
| def cache_key(token_ids): |
| ids = [int(t) for t in token_ids] |
| return hashlib.sha1(struct.pack(f"<{len(ids)}i", *ids)).hexdigest() |
|
|
|
|
| def load_tokenizer(snapshot=None): |
| from transformers import AutoTokenizer |
|
|
| if snapshot is None: |
| from models.autoports.cloudflare_clef.tt.loader import resolve_snapshot |
|
|
| snapshot = resolve_snapshot() |
| return AutoTokenizer.from_pretrained(snapshot) |
|
|
|
|
| def load_processor(snapshot=None): |
| from transformers import AutoProcessor |
|
|
| if snapshot is None: |
| from models.autoports.cloudflare_clef.tt.loader import resolve_snapshot |
|
|
| snapshot = resolve_snapshot() |
| return AutoProcessor.from_pretrained(snapshot) |
|
|
