Skip to content
This repository was archived by the owner on Jul 4, 2025. It is now read-only.

Commit 084592b

Browse files
namchuainamchuai
committed
feat: return model id when download model success
1 parent ba5cd81 commit 084592b

File tree

2 files changed

+112
-27
lines changed

2 files changed

+112
-27
lines changed

‎engine/services/model_service.cc‎

Lines changed: 97 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -2,15 +2,18 @@
22
#include <filesystem>
33
#include <iostream>
44
#include <ostream>
5+
#include "config/gguf_parser.h"
6+
#include "config/yaml_config.h"
57
#include "utils/cli_selection_utils.h"
68
#include "utils/cortexso_parser.h"
79
#include "utils/file_manager_utils.h"
810
#include "utils/huggingface_utils.h"
911
#include "utils/logging_utils.h"
10-
#include "utils/model_callback_utils.h"
12+
#include "utils/modellist_utils.h"
1113
#include "utils/string_utils.h"
1214

13-
void ModelService::DownloadModel(const std::string& input) {
15+
std::optional<std::string> ModelService::DownloadModel(
16+
const std::string& input) {
1417
if (input.empty()) {
1518
throw std::runtime_error(
1619
"Input must be Cortex Model Hub handle or HuggingFace url!");
@@ -32,15 +35,15 @@ void ModelService::DownloadModel(const std::string& input) {
3235
return DownloadModelByModelName(model_name);
3336
}
3437

35-
DownloadHuggingFaceGgufModel(author, model_name, std::nullopt);
3638
CLI_LOG("Model " << model_name << " downloaded successfully!")
37-
return;
39+
return DownloadHuggingFaceGgufModel(author, model_name, std::nullopt);
3840
}
3941

4042
return DownloadModelByModelName(input);
4143
}
4244

43-
void ModelService::DownloadModelByModelName(const std::string& modelName) {
45+
std::optional<std::string> ModelService::DownloadModelByModelName(
46+
const std::string& modelName) {
4447
try {
4548
auto branches =
4649
huggingface_utils::GetModelRepositoryBranches("cortexso", modelName);
@@ -52,12 +55,13 @@ void ModelService::DownloadModelByModelName(const std::string& modelName) {
5255
}
5356
if (options.empty()) {
5457
CLI_LOG("No variant found");
55-
return;
58+
return std::nullopt;
5659
}
5760
auto selection = cli_selection_utils::PrintSelection(options);
58-
DownloadModelFromCortexso(modelName, selection.value());
61+
return DownloadModelFromCortexso(modelName, selection.value());
5962
} catch (const std::runtime_error& e) {
6063
CLI_LOG("Error downloading model, " << e.what());
64+
return std::nullopt;
6165
}
6266
}
6367

@@ -87,7 +91,8 @@ std::optional<config::ModelConfig> ModelService::GetDownloadedModel(
8791
return std::nullopt;
8892
}
8993

90-
void ModelService::DownloadModelByDirectUrl(const std::string& url) {
94+
std::optional<std::string> ModelService::DownloadModelByDirectUrl(
95+
const std::string& url) {
9196
auto url_obj = url_parser::FromUrlString(url);
9297

9398
if (url_obj.host == kHuggingFaceHost) {
@@ -103,6 +108,9 @@ void ModelService::DownloadModelByDirectUrl(const std::string& url) {
103108
return DownloadModelFromCortexso(model_id);
104109
}
105110

111+
std::string huggingFaceHost{kHuggingFaceHost};
112+
std::string unique_model_id{huggingFaceHost + "/" + author + "/" + model_id +
113+
"/" + file_name};
106114
auto local_path{file_manager_utils::GetModelsContainerPath() /
107115
"huggingface.co" / author / model_id / file_name};
108116

@@ -119,41 +127,76 @@ void ModelService::DownloadModelByDirectUrl(const std::string& url) {
119127
auto downloadTask{DownloadTask{.id = model_id,
120128
.type = DownloadType::Model,
121129
.items = {DownloadItem{
122-
.id = url_obj.pathParams.back(),
130+
.id = unique_model_id,
123131
.downloadUrl = download_url,
124132
.localPath = local_path,
125133
}}}};
126134

127-
auto on_finished = [&author](const DownloadTask& finishedTask) {
135+
auto on_finished = [&](const DownloadTask& finishedTask) {
128136
CLI_LOG("Model " << finishedTask.id << " downloaded successfully!")
129137
auto gguf_download_item = finishedTask.items[0];
130-
model_callback_utils::ParseGguf(gguf_download_item, author);
138+
ParseGguf(gguf_download_item, author);
131139
};
132140

133141
download_service_.AddDownloadTask(downloadTask, on_finished);
142+
return unique_model_id;
134143
}
135144

136-
void ModelService::DownloadModelFromCortexso(const std::string& name,
137-
const std::string& branch) {
145+
std::optional<std::string> ModelService::DownloadModelFromCortexso(
146+
const std::string& name, const std::string& branch) {
147+
138148
auto downloadTask = cortexso_parser::getDownloadTask(name, branch);
139149
if (downloadTask.has_value()) {
140-
DownloadService().AddDownloadTask(downloadTask.value(),
141-
model_callback_utils::DownloadModelCb);
142-
CLI_LOG("Model " << name << " downloaded successfully!")
150+
std::string model_id{name + ":" + branch};
151+
DownloadService().AddDownloadTask(
152+
downloadTask.value(), [&](const DownloadTask& finishedTask) {
153+
const DownloadItem* model_yml_item = nullptr;
154+
auto need_parse_gguf = true;
155+
156+
for (const auto& item : finishedTask.items) {
157+
if (item.localPath.filename().string() == "model.yml") {
158+
model_yml_item = &item;
159+
}
160+
}
161+
162+
if (model_yml_item != nullptr) {
163+
auto url_obj =
164+
url_parser::FromUrlString(model_yml_item->downloadUrl);
165+
CTL_INF("Adding model to modellist with branch: " << branch);
166+
config::YamlHandler yaml_handler;
167+
yaml_handler.ModelConfigFromFile(
168+
model_yml_item->localPath.string());
169+
auto mc = yaml_handler.GetModelConfig();
170+
171+
modellist_utils::ModelListUtils modellist_utils_obj;
172+
modellist_utils::ModelEntry model_entry{
173+
.model_id = model_id,
174+
.author_repo_id = "cortexso",
175+
.branch_name = branch,
176+
.path_to_model_yaml = model_yml_item->localPath.string(),
177+
.model_alias = model_id,
178+
.status = modellist_utils::ModelStatus::READY};
179+
modellist_utils_obj.AddModelEntry(model_entry);
180+
}
181+
});
182+
183+
CLI_LOG("Model " << model_id << " downloaded successfully!")
184+
return model_id;
143185
} else {
144186
CTL_ERR("Model not found");
187+
return std::nullopt;
145188
}
146189
}
147190

148-
void ModelService::DownloadHuggingFaceGgufModel(
191+
std::optional<std::string> ModelService::DownloadHuggingFaceGgufModel(
149192
const std::string& author, const std::string& modelName,
150193
std::optional<std::string> fileName) {
151194
auto repo_info =
152195
huggingface_utils::GetHuggingFaceModelRepoInfo(author, modelName);
153196
if (!repo_info.has_value()) {
154197
// throw is better?
155198
CTL_ERR("Model not found");
156-
return;
199+
return std::nullopt;
157200
}
158201

159202
if (!repo_info->gguf.has_value()) {
@@ -172,5 +215,40 @@ void ModelService::DownloadHuggingFaceGgufModel(
172215

173216
auto download_url = huggingface_utils::GetDownloadableUrl(author, modelName,
174217
selection.value());
175-
DownloadModelByDirectUrl(download_url);
218+
return DownloadModelByDirectUrl(download_url);
219+
}
220+
221+
void ModelService::ParseGguf(const DownloadItem& ggufDownloadItem,
222+
std::optional<std::string> author) const {
223+
224+
config::GGUFHandler gguf_handler;
225+
config::YamlHandler yaml_handler;
226+
gguf_handler.Parse(ggufDownloadItem.localPath.string());
227+
config::ModelConfig model_config = gguf_handler.GetModelConfig();
228+
model_config.id =
229+
ggufDownloadItem.localPath.parent_path().filename().string();
230+
model_config.files = {ggufDownloadItem.localPath.string()};
231+
yaml_handler.UpdateModelConfig(model_config);
232+
233+
auto yaml_path{ggufDownloadItem.localPath};
234+
auto yaml_name = yaml_path.replace_extension(".yml");
235+
236+
if (!std::filesystem::exists(yaml_path)) {
237+
yaml_handler.WriteYamlFile(yaml_path.string());
238+
}
239+
240+
auto url_obj = url_parser::FromUrlString(ggufDownloadItem.downloadUrl);
241+
auto branch = url_obj.pathParams[3];
242+
CTL_INF("Adding model to modellist with branch: " << branch);
243+
244+
auto author_id = author.has_value() ? author.value() : "cortexso";
245+
modellist_utils::ModelListUtils modellist_utils_obj;
246+
modellist_utils::ModelEntry model_entry{
247+
.model_id = ggufDownloadItem.id,
248+
.author_repo_id = author_id,
249+
.branch_name = branch,
250+
.path_to_model_yaml = yaml_name.string(),
251+
.model_alias = ggufDownloadItem.id,
252+
.status = modellist_utils::ModelStatus::READY};
253+
modellist_utils_obj.AddModelEntry(model_entry, true);
176254
}

‎engine/services/model_service.h‎

Lines changed: 15 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -8,27 +8,34 @@ class ModelService {
88
public:
99
ModelService() : download_service_{DownloadService()} {};
1010

11-
void DownloadModel(const std::string& input);
11+
/**
12+
* Return model id if download successfully
13+
*/
14+
std::optional<std::string> DownloadModel(const std::string& input);
1215

1316
std::optional<config::ModelConfig> GetDownloadedModel(
1417
const std::string& modelId) const;
1518

1619
private:
17-
void DownloadModelByDirectUrl(const std::string& url);
20+
std::optional<std::string> DownloadModelByDirectUrl(const std::string& url);
1821

19-
void DownloadModelFromCortexso(const std::string& name,
20-
const std::string& branch = "main");
22+
std::optional<std::string> DownloadModelFromCortexso(
23+
const std::string& name, const std::string& branch = "main");
2124

2225
/**
2326
* Handle downloading model which have following pattern: author/model_name
2427
*/
25-
void DownloadHuggingFaceGgufModel(const std::string& author,
26-
const std::string& modelName,
27-
std::optional<std::string> fileName);
28+
std::optional<std::string> DownloadHuggingFaceGgufModel(
29+
const std::string& author, const std::string& modelName,
30+
std::optional<std::string> fileName);
2831

29-
void DownloadModelByModelName(const std::string& modelName);
32+
std::optional<std::string> DownloadModelByModelName(
33+
const std::string& modelName);
3034

3135
DownloadService download_service_;
3236

37+
void ParseGguf(const DownloadItem& ggufDownloadItem,
38+
std::optional<std::string> author = nullptr) const;
39+
3340
constexpr auto static kHuggingFaceHost = "huggingface.co";
3441
};

0 commit comments

Comments
 (0)