add safetensors version

This commit is contained in:
Patrick Devine
2024-04-24 18:32:01 -07:00
committed by Michael Yang
parent d88582dffd
commit 4730762e5c
2 changed files with 20 additions and 4 deletions

View File

@@ -281,6 +281,15 @@ func (m *SafetensorFormat) GetModelArch(name, dirPath string, params *Params) (M
return nil, fmt.Errorf("No architecture specified to convert")
case 1:
switch params.Architectures[0] {
case "LlamaForCausalLM":
return &LlamaModel{
ModelData{
Name: name,
Path: dirPath,
Params: params,
Format: m,
},
}, nil
case "MistralForCausalLM":
return &MistralModel{
ModelData{