add rm command for models (#151)

This commit is contained in:
Patrick Devine
2023-07-20 16:09:23 -07:00
committed by GitHub
parent 8454f298ac
commit e7a393de54
5 changed files with 166 additions and 25 deletions

View File

@@ -487,6 +487,83 @@ func CreateLayer(f io.ReadSeeker) (*LayerReader, error) {
return layer, nil
}
func DeleteModel(name string, fn func(api.ProgressResponse)) error {
mp := ParseModelPath(name)
manifest, err := GetManifest(mp)
if err != nil {
fn(api.ProgressResponse{Status: "couldn't retrieve manifest"})
return err
}
deleteMap := make(map[string]bool)
for _, layer := range manifest.Layers {
deleteMap[layer.Digest] = true
}
deleteMap[manifest.Config.Digest] = true
fp, err := GetManifestPath()
if err != nil {
fn(api.ProgressResponse{Status: "problem getting manifest path"})
return err
}
err = filepath.Walk(fp, func(path string, info os.FileInfo, err error) error {
if err != nil {
fn(api.ProgressResponse{Status: "problem walking manifest dir"})
return err
}
if !info.IsDir() {
path := path[len(fp)+1:]
slashIndex := strings.LastIndex(path, "/")
if slashIndex == -1 {
return nil
}
tag := path[:slashIndex] + ":" + path[slashIndex+1:]
fmp := ParseModelPath(tag)
// skip the manifest we're trying to delete
if mp.GetFullTagname() == fmp.GetFullTagname() {
return nil
}
// save (i.e. delete from the deleteMap) any files used in other manifests
manifest, err := GetManifest(fmp)
if err != nil {
log.Printf("skipping file: %s", fp)
return nil
}
for _, layer := range manifest.Layers {
delete(deleteMap, layer.Digest)
}
delete(deleteMap, manifest.Config.Digest)
}
return nil
})
// only delete the files which are still in the deleteMap
for k, v := range deleteMap {
if v {
err := os.Remove(k)
if err != nil {
log.Printf("couldn't remove file '%s': %v", k, err)
continue
}
}
}
fp, err = mp.GetManifestPath(false)
if err != nil {
return err
}
err = os.Remove(fp)
if err != nil {
log.Printf("couldn't remove manifest file '%s': %v", fp, err)
return err
}
fn(api.ProgressResponse{Status: fmt.Sprintf("deleted '%s'", name)})
return nil
}
func PushModel(name, username, password string, fn func(api.ProgressResponse)) error {
mp := ParseModelPath(name)

View File

@@ -18,7 +18,7 @@ import (
"github.com/jmorganca/ollama/llama"
)
func generate(c *gin.Context) {
func GenerateHandler(c *gin.Context) {
start := time.Now()
var req api.GenerateRequest
@@ -78,7 +78,7 @@ func generate(c *gin.Context) {
streamResponse(c, ch)
}
func pull(c *gin.Context) {
func PullModelHandler(c *gin.Context) {
var req api.PullRequest
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
@@ -100,7 +100,7 @@ func pull(c *gin.Context) {
streamResponse(c, ch)
}
func push(c *gin.Context) {
func PushModelHandler(c *gin.Context) {
var req api.PushRequest
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
@@ -122,7 +122,7 @@ func push(c *gin.Context) {
streamResponse(c, ch)
}
func create(c *gin.Context) {
func CreateModelHandler(c *gin.Context) {
var req api.CreateRequest
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"message": err.Error()})
@@ -146,7 +146,30 @@ func create(c *gin.Context) {
streamResponse(c, ch)
}
func list(c *gin.Context) {
func DeleteModelHandler(c *gin.Context) {
var req api.DeleteRequest
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
ch := make(chan any)
go func() {
defer close(ch)
fn := func(r api.ProgressResponse) {
ch <- r
}
if err := DeleteModel(req.Name, fn); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
}()
streamResponse(c, ch)
}
func ListModelsHandler(c *gin.Context) {
var models []api.ListResponseModel
fp, err := GetManifestPath()
if err != nil {
@@ -199,11 +222,12 @@ func Serve(ln net.Listener) error {
c.String(http.StatusOK, "Ollama is running")
})
r.POST("/api/pull", pull)
r.POST("/api/generate", generate)
r.POST("/api/create", create)
r.POST("/api/push", push)
r.GET("/api/tags", list)
r.POST("/api/pull", PullModelHandler)
r.POST("/api/generate", GenerateHandler)
r.POST("/api/create", CreateModelHandler)
r.POST("/api/push", PushModelHandler)
r.GET("/api/tags", ListModelsHandler)
r.DELETE("/api/delete", DeleteModelHandler)
log.Printf("Listening on %s", ln.Addr())
s := &http.Server{