Small fixup to template loading

2 years ago · 4a932483e1
parent b710147b95
commit 4a932483e1
1 changed files with 2 additions and 1 deletions
--- a/model_loader.go
+++ b/model_loader.go
@ -62,6 +62,7 @@ func (ml *ModelLoader) LoadModel(modelName string, opts ...llama.ModelOption) (*
 		if _, err := os.Stat(modelBin); os.IsNotExist(err) {
 			return nil, err
 		} else {
+			modelName = fmt.Sprintf("%s.bin", modelName)
 			modelFile = modelBin
 		}
 	}
@ -87,7 +88,7 @@ func (ml *ModelLoader) LoadModel(modelName string, opts ...llama.ModelOption) (*
 		if err != nil {
 			return nil, err
 		}
-		ml.promptsTemplates[modelFile] = tmpl
+		ml.promptsTemplates[modelName] = tmpl
 	}

 	ml.models[modelFile] = model