From 28adbf5e24ef6f96a8db51904d818fa9d9455b73 Mon Sep 17 00:00:00 2001
From: zouap <zouap@pcl.ac.cn>
Date: Fri, 31 Dec 2021 10:26:07 +0800
Subject: [PATCH] =?UTF-8?q?=E6=8F=90=E4=BA=A4=E4=BB=A3=E7=A0=81=E3=80=82?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: zouap <zouap@pcl.ac.cn>
---
 routers/repo/ai_model_manage.go | 50 +++++++++++++++++++++++++++++++++
 routers/routes/routes.go        |  2 ++
 2 files changed, 52 insertions(+)

diff --git a/routers/repo/ai_model_manage.go b/routers/repo/ai_model_manage.go
index 669bdf9fa..86d8c140a 100644
--- a/routers/repo/ai_model_manage.go
+++ b/routers/repo/ai_model_manage.go
@@ -586,3 +586,53 @@ func ModifyModelInfo(ctx *context.Context) {
 	}
 
 }
+
+func QueryModelListForPredict(ctx *context.Context) {
+	repoId := ctx.Repo.Repository.ID
+	modelResult, count, err := models.QueryModel(&models.AiModelQueryOptions{
+		ListOptions: models.ListOptions{
+			Page:     -1,
+			PageSize: -1,
+		},
+		RepoID: repoId,
+		Type:   -1,
+		New:    -1,
+	})
+	if err != nil {
+		ctx.ServerError("Cloudbrain", err)
+		return
+	}
+	log.Info("query return count=" + fmt.Sprint(count))
+
+	nameList := make([]string, 0)
+
+	nameMap := make(map[string][]*models.AiModelManage)
+	for _, model := range modelResult {
+		if _, value := nameMap[model.Name]; !value {
+			models := make([]*models.AiModelManage, 0)
+			models = append(models, model)
+			nameMap[model.Name] = models
+			nameList = append(nameList, model.Name)
+		} else {
+			nameMap[model.Name] = append(nameMap[model.Name], model)
+		}
+	}
+
+	mapInterface := make(map[string]interface{})
+	mapInterface["nameList"] = nameList
+	mapInterface["nameMap"] = nameMap
+	ctx.JSON(http.StatusOK, mapInterface)
+}
+
+func QueryModelFileForPredict(ctx *context.Context) {
+	id := ctx.Query("ID")
+	model, err := models.QueryModelById(id)
+	if err != nil {
+		log.Error("no such model!", err.Error())
+		ctx.ServerError("no such model:", err)
+		return
+	}
+	prefix := model.Path[len(setting.Bucket)+1:]
+	fileinfos, err := storage.GetAllObjectByBucketAndPrefix(setting.Bucket, prefix)
+	ctx.JSON(http.StatusOK, fileinfos)
+}
diff --git a/routers/routes/routes.go b/routers/routes/routes.go
index 481b0eacb..547e48a4c 100755
--- a/routers/routes/routes.go
+++ b/routers/routes/routes.go
@@ -992,6 +992,8 @@ func RegisterRoutes(m *macaron.Macaron) {
 			m.Get("/show_model_child_api", repo.ShowOneVersionOtherModel)
 			m.Get("/query_train_job", reqRepoCloudBrainReader, repo.QueryTrainJobList)
 			m.Get("/query_train_job_version", reqRepoCloudBrainReader, repo.QueryTrainJobVersionList)
+			m.Get("/query_model_for_predict", reqRepoCloudBrainReader, repo.QueryModelListForPredict)
+			m.Get("/query_modelfile_for_predict", reqRepoCloudBrainReader, repo.QueryModelFileForPredict)
 			m.Group("/:ID", func() {
 				m.Get("", repo.ShowSingleModel)
 				m.Get("/downloadsingle", repo.DownloadSingleModelFile)