Skip to content

Commit d257c50

Browse files
authored
Merge pull request #626 from yuerqiqi/my-feature-dev
feat: implement qwen3 probing service and tests
2 parents 4fd29d0 + 8b49bad commit d257c50

6 files changed

Lines changed: 1878 additions & 4 deletions

File tree

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,16 @@
11
add_executable(mllm-qwen3-service main.cpp)
22
target_link_libraries(mllm-qwen3-service PRIVATE MllmRT MllmCPUBackend)
33
target_include_directories(mllm-qwen3-service PRIVATE ${MLLM_INCLUDE_DIR})
4+
5+
add_executable(mllm-qwen3-accuracy test_accuracy.cpp)
6+
target_link_libraries(mllm-qwen3-accuracy PRIVATE MllmRT MllmCPUBackend)
7+
target_include_directories(mllm-qwen3-accuracy PRIVATE ${MLLM_INCLUDE_DIR})
8+
9+
10+
add_executable(mllm-qwen3-probing main_probing.cpp)
11+
target_link_libraries(mllm-qwen3-probing PRIVATE MllmRT MllmCPUBackend)
12+
target_include_directories(mllm-qwen3-probing PRIVATE ${MLLM_INCLUDE_DIR})
13+
14+
add_executable(mllm-qwen3-trivia-probing test_trivia_probing.cpp)
15+
target_link_libraries(mllm-qwen3-trivia-probing PRIVATE MllmRT MllmCPUBackend)
16+
target_include_directories(mllm-qwen3-trivia-probing PRIVATE ${MLLM_INCLUDE_DIR})
Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
#include <string>
2+
#include <vector>
3+
#include <filesystem>
4+
#include <iostream>
5+
#include <sstream>
6+
7+
#include <fmt/core.h>
8+
#include <fmt/color.h>
9+
#include <nlohmann/json.hpp>
10+
11+
#include <mllm/mllm.hpp>
12+
#include <mllm/engine/service/Service.hpp>
13+
#include "mllm/models/qwen3/modeling_qwen3_probing_service.hpp"
14+
15+
using namespace mllm;
16+
using namespace mllm::models::qwen3_probing;
17+
namespace fs = std::filesystem;
18+
19+
std::vector<int> parseLayers(const std::string& input) {
20+
std::vector<int> layers;
21+
std::stringstream ss(input);
22+
std::string segment;
23+
int num;
24+
ss >> segment;
25+
while (ss >> num) layers.push_back(num);
26+
return layers;
27+
}
28+
29+
MLLM_MAIN({
30+
mllm::setLogLevel(mllm::LogLevel::kError);
31+
auto& model_path = mllm::Argparse::add<std::string>("-m|--model_path").help("Model path").required(true);
32+
auto& probe_path = mllm::Argparse::add<std::string>("-p|--probe_path").help("Probes dir").required(true);
33+
mllm::Argparse::parse(argc, argv);
34+
35+
auto qwen3_session = std::make_shared<Qwen3ProbingSession>();
36+
try {
37+
qwen3_session->fromPreTrain(model_path.get());
38+
} catch (const std::exception& e) {
39+
std::cerr << "Load Model Error: " << e.what() << std::endl;
40+
return 1;
41+
}
42+
43+
ProbingArgs p_args;
44+
p_args.enable_prefill_check = true;
45+
p_args.enable_decode_check = true;
46+
47+
p_args.prefill_stop_threshold = 0.7f;
48+
p_args.decode_stop_threshold = 0.8f;
49+
50+
p_args.pos_threshold = 0.9f;
51+
52+
std::cout << ">>> Loading Probes..." << std::endl;
53+
qwen3_session->setProbingArgs(p_args);
54+
qwen3_session->loadProbes(probe_path.get(), p_args);
55+
56+
mllm::service::insertSession("mllmTeam/Qwen3-Probing", qwen3_session);
57+
mllm::service::startService();
58+
59+
std::vector<nlohmann::json> history;
60+
std::vector<int> current_prefill_layers = {27, 30}; // 默认
61+
62+
std::cout << "\n[System] Ready. Commands:\n";
63+
std::cout << " /prefill 15 20 Set prefill layers\n";
64+
std::cout << " /clear Clear history\n";
65+
std::cout << " /exit Exit\n";
66+
67+
while (true) {
68+
std::cout << "\nUser: ";
69+
std::string user_input;
70+
std::getline(std::cin, user_input);
71+
72+
if (user_input == "/exit") break;
73+
if (user_input == "/clear") {
74+
history.clear();
75+
continue;
76+
}
77+
if (user_input.rfind("/prefill", 0) == 0) {
78+
current_prefill_layers = parseLayers(user_input);
79+
std::cout << "Prefill layers: " << nlohmann::json(current_prefill_layers).dump() << "\n";
80+
continue;
81+
}
82+
83+
nlohmann::json user_msg;
84+
user_msg["role"] = "user";
85+
user_msg["content"] = user_input;
86+
history.push_back(user_msg);
87+
88+
nlohmann::json req;
89+
req["model"] = "mllmTeam/Qwen3-Probing";
90+
req["messages"] = history;
91+
req["prefill_layers"] = current_prefill_layers;
92+
req["enable_thinking"] = false;
93+
req["id"] = "chat-probing";
94+
95+
mllm::service::sendRequest(req.dump());
96+
97+
std::string assistant_content;
98+
bool thinking = false;
99+
100+
while (true) {
101+
std::string resp = mllm::service::getResponse("chat-probing");
102+
auto j = nlohmann::json::parse(resp);
103+
104+
if (j.contains("choices") && j["choices"].size() > 0) {
105+
auto& choice = j["choices"][0];
106+
auto content = choice["delta"]["content"];
107+
108+
if (content.is_string()) {
109+
std::string s = content.get<std::string>();
110+
if (s.find("early_exit") != std::string::npos) {
111+
try {
112+
auto warn = nlohmann::json::parse(s);
113+
fmt::print(fmt::fg(fmt::color::red) | fmt::emphasis::bold,
114+
"\n[Hallucination] Phase: {} | Layer: {} | Score: {:.4f}\n", warn.value("phase", "unknown"),
115+
warn.value("layer", -1), warn.value("score", 0.0f));
116+
} catch (...) { fmt::print(fmt::fg(fmt::color::red), "\n[Hallucination] Raw: {}\n", s); }
117+
118+
if (!history.empty() && history.back()["role"] == "user") history.pop_back();
119+
break;
120+
}
121+
122+
if (s == "<think>") {
123+
thinking = true;
124+
continue;
125+
}
126+
if (s == "</think>") {
127+
thinking = false;
128+
continue;
129+
}
130+
131+
if (thinking)
132+
fmt::print(fmt::fg(fmt::color::gray), "{}", s);
133+
else {
134+
fmt::print("{}", s);
135+
assistant_content += s;
136+
}
137+
std::fflush(stdout);
138+
}
139+
140+
if (choice["finish_reason"] == "stop") break;
141+
}
142+
}
143+
144+
if (!assistant_content.empty()) {
145+
nlohmann::json msg;
146+
msg["role"] = "assistant";
147+
msg["content"] = assistant_content;
148+
history.push_back(msg);
149+
}
150+
}
151+
mllm::service::stopService();
152+
return 0;
153+
})

0 commit comments

Comments
 (0)