return sd

def fix_bf16(weights:dict[Any, Tensor]):

# TODO: without casting to float16, 70B llama OOM on tinybox.

return {k:v.cast(dtypes.float32).cast(dtypes.float16) if v.dtype == dtypes.bfloat16 else v for k,v in weights.items()}

class Tokenizer:

pat_str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"

def __init__(self, model_path: str):

# load weights

weights = load(str(model_path / "model.safetensors"))

weights = convert_from_huggingface(weights, MODEL_PARAMS[model_size]["args"]["n_layers"], MODEL_PARAMS[model_size]["args"]["n_heads"], MODEL_PARAMS[model_size]["args"]["n_kv_heads"])

weights = fix_bf16(weights)

with Context(BEAM=0):

# shard