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 << " \n User: " ;
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