tokenizer.py

from pathlib import Path
from transformers import AutoTokenizer
import mlx.core as mx

MODEL = str(Path(__file__).parent / "llama3.2/snapshots/4e20de362430cd3b72f300e6b0f18e50e7166e08")

SYSTEM_PROMPT = (
    "You are Lancer, a concise and helpful AI assistant. "
    "Answer clearly and accurately, and say so when you are unsure."
)

def build_chat_prompt(user_text, system_prompt=SYSTEM_PROMPT):
    return (
        "<|begin_of_text|>"
        "<|start_header_id|>system<|end_header_id|>\n\n"
        f"{system_prompt}<|eot_id|>"
        "<|start_header_id|>user<|end_header_id|>\n\n"
        f"{user_text}<|eot_id|>"
        "<|start_header_id|>assistant<|end_header_id|>\n\n"
    )

def encode(text):
    tokenizer = AutoTokenizer.from_pretrained(MODEL)
    prompt = build_chat_prompt(text)
    tokens_ids = tokenizer.encode(prompt, add_special_tokens=False)
    return mx.array(tokens_ids)
     

def get_input_embeds(tokens_ids):
    weights = mx.load(f"{MODEL}/model.safetensors")
    w = weights["model.embed_tokens.weight"]
    embeds = w[tokens_ids]
    return embeds


def decode(token_ids):
    tokenizer = AutoTokenizer.from_pretrained(MODEL)
    return tokenizer.decode(token_ids)