try:

exceptImportError:

print("Please install `safetensors` python package")

sys.exit(1)

defcount_model_parts(dir_model:Path)->int:

forfilenameinos.listdir(dir_model):

iffilename.startswith("model-00"):

ifnum_parts>0:

print("gguf: found "+str(num_parts)+" model parts")

defparse_args()->argparse.Namespace:

parser=argparse.ArgumentParser(description="Convert a PLaMo model to a GGML compatible file")

parser.add_argument(

"--vocab-only",action="store_true",

help="extract only the vocab",

)

parser.add_argument(

"--outfile",type=Path,

help="path to write to; default: based on input",

)

parser.add_argument(

"model",type=Path,

help="directory containing model file, or model file itself (*.bin)",

)

parser.add_argument(

"ftype",type=int,choices=[0,1],default=1,nargs='?',

help="output format - use 0 for float32, 1 for float16",

)

returnparser.parse_args()

args=parse_args()

dir_model=args.model

ftype=args.ftype

ifnotdir_model.is_dir():

print(f'Error:{args.model} is not a directory',file=sys.stderr)

sys.exit(1)

ftype_str= ["f32","f16"]

ifargs.outfileisnotNone:

fname_out=args.outfile

else:

print("gguf: loading model "+dir_model.name)

withopen(dir_model/"config.json","r",encoding="utf-8")asf:

hparams=json.load(f)

ifhparams["architectures"][0]!="PlamoForCausalLM":

print("Model architecture not supported: "+hparams["architectures"][0])

sys.exit(1)

num_parts=count_model_parts(dir_model)

ARCH=gguf.MODEL_ARCH.PLAMO

gguf_writer=gguf.GGUFWriter(fname_out,gguf.MODEL_ARCH_NAMES[ARCH])

print("gguf: get model metadata")

block_count=hparams["num_hidden_layers"]

gguf_writer.add_name("PLaMo")

gguf_writer.add_context_length(4096)# not in config.json

gguf_writer.add_embedding_length(hparams["hidden_size"])

gguf_writer.add_feed_forward_length(hparams["intermediate_size"])

gguf_writer.add_block_count(block_count)

gguf_writer.add_head_count(hparams["num_attention_heads"])

gguf_writer.add_head_count_kv(hparams["num_attention_heads"]//hparams["n_shared_head"])

gguf_writer.add_layer_norm_rms_eps(hparams["rms_norm_eps"])

gguf_writer.add_file_type(ftype)

print("gguf: get tokenizer metadata")

tokens:list[bytes]= []

scores:list[float]= []

toktypes:list[int]= []

ifnottokenizer_model_file.is_file():

print(f'Error: Missing{tokenizer_model_file}',file=sys.stderr)

sys.exit(1)

print("gguf: get sentencepiece tokenizer vocab, scores and token types")

tokenizer=SentencePieceProcessor(str(tokenizer_model_file))

foriinrange(tokenizer.vocab_size()):

text:bytes

score:float

piece=tokenizer.id_to_piece(i)

text=piece.encode("utf-8")

score=tokenizer.get_score(i)

iftokenizer.is_unknown(i):

iftokenizer.is_control(i):

iftokenizer.is_unused(i):

iftokenizer.is_byte(i):

tokens.append(text)

scores.append(score)

toktypes.append(toktype)

gguf_writer.add_tokenizer_model("llama")

gguf_writer.add_token_list(tokens)

gguf_writer.add_token_scores(scores)

gguf_writer.add_token_types(toktypes)

gguf_writer.add_sep_token_id(5)

gguf_writer.add_pad_token_id(3)

special_vocab=gguf.SpecialVocab(dir_model)

special_vocab.add_to_gguf(gguf_writer)

tensor_map=gguf.get_tensor_name_map(ARCH,block_count)

n_head=hparams["num_attention_heads"]

n_head_kv=hparams["num_key_value_heads"]

head_dim=hparams["hidden_size"]//n_head

print("gguf: get tensor metadata")

ifnum_parts==0:

part_names=iter(("model.safetensors",))

else:

part_names= (

f"model-{n:05}-of-{num_parts:05}.safetensors"forninrange(1,num_parts+1)

)

forpart_nameinpart_names:

ifargs.vocab_only:

print("gguf: loading model part '"+part_name+"'")

model_part=safe_open(dir_model/part_name,framework="pt")

fornameinmodel_part.keys():

if"self_attn.rotary_emb.inv_freq"inname:

data=model_part.get_tensor(name)

old_dtype=data.dtype

ifdata.dtype!=torch.float16anddata.dtype!=torch.float32:

data=data.to(torch.float32)

data=data.squeeze().numpy()

new_name=tensor_map.get_name(name,try_suffixes= (".weight",".bias"))

ifnew_nameisNone:

print("Can not map tensor '"+name+"'")

sys.exit()

n_dims=len(data.shape)

data_dtype=data.dtype

ifftype==0anddata_dtype==np.float16:

data=data.astype(np.float32)

ifftype==1anddata_dtype==np.float16andn_dims==1:

data=data.astype(np.float32)

ifftype==1anddata_dtype==np.float32andname.endswith(".weight")andn_dims==2:

data=data.astype(np.float16)

print(new_name+", n_dims = "+str(n_dims)+", "+str(old_dtype)+" --> "+str(data.dtype))

gguf_writer.add_tensor(new_name,data)

print("gguf: write header")

gguf_writer.write_header_to_file()

print("gguf: write metadata")

gguf_writer.write_kv_data_to_file()

ifnotargs.vocab_only:

print("gguf: write tensors")

gguf_writer.write_tensors_to_file()

gguf_writer.close()

print(f"gguf: model successfully exported to '{fname_out}'")

print("")

ggml_element_size(kv_self.k)*n_embd_head,

ggml_element_size(kv_self.k)*n_embd_gqa,

ggml_element_size(kv_self.k)*n_embd_gqa*n_ctx*il),

K_shift, n_embd_head,0,0, freq_base, freq_scale);

K_shift, n_embd_head,2,0, freq_base, freq_scale);

offload_func_kq(tmp);

ggml_build_forward_expand(gf, tmp);

}

offload_func_kq(tmpq);

ggml_set_name(tmpq,"tmpq");

structggml_tensor * Kcur =ggml_rope_custom(ctx0,ggml_reshape_3d(ctx0, tmpk, n_embd_head, n_head_kv, n_tokens), KQ_pos, n_embd_head,0,0, freq_base, freq_scale);

structggml_tensor * Kcur =ggml_rope_custom(ctx0,ggml_reshape_3d(ctx0, tmpk, n_embd_head, n_head_kv, n_tokens), KQ_pos, n_embd_head,2,0, freq_base, freq_scale);

offload_func_kq(Kcur);

ggml_set_name(Kcur,"Kcur");

structggml_tensor * Qcur =ggml_rope_custom(ctx0,ggml_reshape_3d(ctx0, tmpq, n_embd_head, n_head, n_tokens), KQ_pos, n_embd_head,0,0, freq_base, freq_scale);

structggml_tensor * Qcur =ggml_rope_custom(ctx0,ggml_reshape_3d(ctx0, tmpq, n_embd_head, n_head, n_tokens), KQ_pos, n_embd_head,2,0, freq_base, freq_scale);

offload_func_kq(Qcur);

ggml_set_name(Qcur,"Qcur");

offload_func_kq(K);

ggml_set_name(K,"K");

structggml_tensor * KQ =ggml_mul_mat(ctx0, K, Q);

structggml_tensor * K_repeated =ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, K->ne[0], K->ne[1], Q->ne[2]);

offload_func_kq(K_repeated);

ggml_set_name(K_repeated,"K_repeated");

structggml_tensor * KQ =ggml_mul_mat(ctx0,ggml_repeat(ctx0, K, K_repeated), Q);

offload_func_kq(KQ);

ggml_set_name(KQ,"KQ");

offload_func_v(V);

ggml_set_name(V,"V");

#if1

structggml_tensor * KQV =ggml_mul_mat(ctx0, V, KQ_soft_max);

structggml_tensor * V_repeated =ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, V->ne[0], V->ne[1], Q->ne[2]);

offload_func_v(V_repeated);

ggml_set_name(V_repeated,"V_repeated");

structggml_tensor * KQV =ggml_mul_mat(ctx0,ggml_repeat(ctx0, V, V_repeated), KQ_soft_max);

offload_func_v(KQV);

ggml_set_name(KQV,"KQV");

#else

#endif

structggml_tensor * KQV_merged =ggml_permute(ctx0, KQV,0,2,1,3);