Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 24 additions & 19 deletions edge0-infer/src/models/edge0_35b/output.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -2,17 +2,18 @@
#include "model/json.h"
#include <set>
#include <stdexcept>
#include <string_view>

namespace edge0::models::edge0_35b {
namespace {
using model::ParsedKind;
using model::ParsedOutput;
using model::ParsedReceiver;
std::size_t character_size(const std::string& text) {
std::size_t character_size(std::string_view text) {
const auto c=static_cast<unsigned char>(text.front());
return c<128?1:c<224?2:c<240?3:4;
}
bool whitespace(const std::string& text) {
bool whitespace(std::string_view text) {
unsigned cp=static_cast<unsigned char>(text.front());
const auto size=character_size(text);
if(size>1) {
Expand All @@ -23,23 +24,24 @@ bool whitespace(const std::string& text) {
cp==0x1680||(cp>=0x2000&&cp<=0x200a)||cp==0x2028||cp==0x2029||
cp==0x202f||cp==0x205f||cp==0x3000;
}
std::string trim(const std::string& text) {
std::string trim(std::string_view text) {
std::size_t first=text.size(), last=0;
for(std::size_t i=0;i<text.size();) {
const auto suffix=text.substr(i);
const auto n=character_size(suffix);
if(!whitespace(suffix)) {if(first==text.size()) first=i;last=i+n;}
i+=n;
}
return first==text.size()?std::string{}:text.substr(first,last-first);
return first==text.size()?std::string{}:std::string(text.substr(first,last-first));
}
class Parser final : public model::OutputParser {
enum class State { reasoning, answer, function, parameter, value, closed, broken };
State state_;
bool tools_, leading_, classification_, string_value_=false, stopped_=false;
std::string buffer_, spaces_, raw_, arguments_, value_;
ToolCall call_;
std::size_t index_=0;
std::size_t index_=0,offset_=0;
std::string_view pending() const {return std::string_view(buffer_).substr(offset_);}
std::set<std::string> keys_;
std::vector<model::Json> definitions_;
static constexpr const char* end_call="</tool_call>";
Expand All @@ -50,14 +52,15 @@ class Parser final : public model::OutputParser {
return true;
}
std::string consume(std::size_t n,bool in_call=true) {
auto text=buffer_.substr(0,n);buffer_.erase(0,n);
auto text=buffer_.substr(offset_,n);offset_+=text.size();
if(in_call) raw_+=text;
return text;
}
bool prefix(const std::string& marker) const {
return buffer_.size()<marker.size() && marker.compare(0,buffer_.size(),buffer_)==0;
const auto text=pending();
return text.size()<marker.size() && marker.compare(0,text.size(),text)==0;
}
bool starts(const std::string& marker) const {return buffer_.compare(0,marker.size(),marker)==0;}
bool starts(const std::string& marker) const {return buffer_.compare(offset_,marker.size(),marker)==0;}
bool argument(const std::string& text,const ParsedReceiver& receive) {
arguments_+=text;return emit(ParsedKind::tool_call_arguments,text,receive);
}
Expand All @@ -81,7 +84,7 @@ class Parser final : public model::OutputParser {
++index_;state_=State::answer;leading_=true;spaces_.clear();return true;
}
bool process(const ParsedReceiver& receive,bool final) {
while(!buffer_.empty() && !stopped_) {
while(!pending().empty() && !stopped_) {
const bool plain=state_==State::answer || state_==State::reasoning;
if(plain) {
if(state_==State::reasoning && starts("</think>")) {
Expand All @@ -93,7 +96,7 @@ class Parser final : public model::OutputParser {
}
if(!final && ((state_==State::reasoning && prefix("</think>")) ||
(tools_ && prefix("<tool_call>")))) break;
const auto n=character_size(buffer_);
const auto n=character_size(pending());
auto text=consume(n,false);
if(whitespace(text)) {
if(!leading_) spaces_+=text;
Expand All @@ -110,7 +113,7 @@ class Parser final : public model::OutputParser {
continue;
}
if(!final && prefix(end_call)) break;
if(state_==State::broken) {consume(character_size(buffer_));continue;}
if(state_==State::broken) {consume(character_size(pending()));continue;}
if(state_==State::value) {
const std::string delimiter="\n</parameter>";
if(starts(delimiter)) {
Expand All @@ -125,14 +128,14 @@ class Parser final : public model::OutputParser {
state_=State::parameter;continue;
}
if(!final && prefix(delimiter)) break;
auto text=consume(character_size(buffer_));
auto text=consume(character_size(pending()));
if(string_value_) {
auto quoted=model::json_quote(text);
if(!argument(quoted.substr(1,quoted.size()-2),receive)) return false;
} else value_+=text;
continue;
}
if(whitespace(buffer_)) {consume(character_size(buffer_));continue;}
if(whitespace(pending())) {consume(character_size(pending()));continue;}
if(state_==State::closed) {state_=State::broken;continue;}
const std::string open=state_==State::function?"<function=":"<parameter=";
if(state_==State::parameter && starts("</function>")) {
Expand All @@ -142,16 +145,17 @@ class Parser final : public model::OutputParser {
}
if(!final && (prefix(open) || (state_==State::parameter && prefix("</function>")))) break;
if(!starts(open)) {state_=State::broken;continue;}
const auto close=buffer_.find('>',open.size());
const auto text=pending();
const auto close=text.find('>',open.size());
if(close==std::string::npos) {if(!final) break;state_=State::broken;continue;}
auto name=trim(buffer_.substr(open.size(),close-open.size()));
auto name=trim(text.substr(open.size(),close-open.size()));
if(name.empty()) {state_=State::broken;continue;}
if(state_==State::function) {
consume(close+1);call_.name=name;state_=State::parameter;
if(!emit(ParsedKind::tool_call_begin,"",receive)) return false;
} else {
if(buffer_.size()==close+1 && !final) break;
if(buffer_.size()==close+1 || buffer_[close+1]!='\n' || keys_.count(name)) {
if(text.size()==close+1 && !final) break;
if(text.size()==close+1 || text[close+1]!='\n' || keys_.count(name)) {
state_=State::broken;continue;
}
consume(close+2);
Expand Down Expand Up @@ -181,15 +185,16 @@ class Parser final : public model::OutputParser {
bool feed(const std::string& text,const ParsedReceiver& receive) override {
if(stopped_) return false;
// Preserve the pre-classification delivery granularity on the legacy path.
if(!tools_ && state_==State::answer && !leading_ && buffer_.empty() && spaces_.empty())
if(!tools_ && state_==State::answer && !leading_ && pending().empty() && spaces_.empty())
return text.empty() || emit(ParsedKind::answer,text,receive);
buffer_.erase(0,offset_);offset_=0;
buffer_+=text;return process(receive,false);
}
bool finish(const ParsedReceiver& receive) override {
if(stopped_) return false;
if(!process(receive,true)) return false;
if(state_!=State::answer && state_!=State::reasoning) {
raw_+=buffer_;buffer_.clear();
raw_+=pending();buffer_.clear();offset_=0;
return emit(ParsedKind::tool_call_rejected,raw_,receive,"incomplete");
}
if(state_==State::answer && !spaces_.empty()) {
Expand Down
46 changes: 46 additions & 0 deletions edge0-infer/tests/unit/models/edge0_35b/output_parser_test.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,36 @@
#include "model/json.h"
#include "tests/support/check.h"
#include <tuple>
#include <algorithm>
#include <cstdlib>
#include <new>
namespace allocation_probe {
bool active=false;
std::size_t maximum=0,large=0;
constexpr std::size_t threshold=32768;
// Keep the deallocator out of line so GCC does not diagnose an inlined free
// against the matching replacement operator new used by this test.
#if defined(__GNUC__)
__attribute__((noinline))
#elif defined(_MSC_VER)
__declspec(noinline)
#endif
void release(void* p) noexcept {std::free(p);}
}
void* operator new(std::size_t bytes) {
if (allocation_probe::active) {
allocation_probe::maximum=std::max(allocation_probe::maximum,bytes);
if (bytes>=allocation_probe::threshold) ++allocation_probe::large;
}
if (auto* p=std::malloc(bytes?bytes:1)) return p;
throw std::bad_alloc();
}
void operator delete(void* p) noexcept {allocation_probe::release(p);}
void operator delete(void* p,std::size_t) noexcept {allocation_probe::release(p);}
void* operator new[](std::size_t bytes) {return ::operator new(bytes);}
void operator delete[](void* p) noexcept {allocation_probe::release(p);}
void operator delete[](void* p,std::size_t) noexcept {allocation_probe::release(p);}

using namespace edge0;
using model::ParsedKind;
using Event = std::tuple<ParsedKind,std::string,std::size_t,std::string,std::string,std::string>;
Expand Down Expand Up @@ -42,6 +72,22 @@ std::string joined(const std::vector<Event>& events,ParsedKind kind) {
}
std::string block(const std::string& body) {return "<tool_call>\n<function=f>\n"+body+"</function>\n</tool_call>";}
int main() {
{
auto parser=models::edge0_35b::output_parser({false,{}, {tool}});
const std::string name(65536,'f');
const auto input="<tool_call><function="+name+"></function></tool_call>";
std::size_t begins=0,ends=0;
const auto receive=[&](const model::ParsedOutput& event) {
if (event.kind==ParsedKind::tool_call_begin) {++begins;CHECK(event.call.name==name);}
if (event.kind==ParsedKind::tool_call_end) {++ends;CHECK(event.call.arguments_json=="{}");}
return true;
};
allocation_probe::large=0;allocation_probe::active=true;
CHECK(parser->feed(input,receive));CHECK(parser->finish(receive));
allocation_probe::active=false;
std::cout << "large_parser_allocations=" << allocation_probe::large << "\n";
CHECK(begins==1 && ends==1 && allocation_probe::large<=32);
}
const auto complete=block("<parameter=s>\nhello\n</parameter>\n<parameter=n>\n3\n</parameter>\n");
const std::vector<std::string> cases={
" \nreason \t</think>\n\nanswer ", "reason"+complete, "<tool_call>literal",
Expand Down