refactor: remove probs from the token vector

Author: pminev

ac-local-plugin/code/LocalLlama.cpp

@@ -347,19 +347,16 @@ struct LocalLlama {
 
         std::vector<int32_t> tokens(tokenData.size());
         std::vector<float> logits(tokenData.size());
-        std::vector<float> probs(tokenData.size());
         for (size_t i = 0; i < tokenData.size(); i++) {
             tokens[i] = tokenData[i].token;
             logits[i] = tokenData[i].logit;
-            probs[i] = tokenData[i].prob;
         }
 
         instance.stopSession();
 
         co_await io.push(Frame_from(sc::StateGeneralInstance::OpGetTokenData{}, {
             .tokens = std::move(tokens),
             .logits = std::move(logits),
-            .probs = std::move(probs)
         }));
     }
 
@@ -368,33 +365,23 @@ struct LocalLlama {
         IoEndpoint& io,
         const sc::StateGeneralInstance::OpCompareTokenData::Params& iparams) {
 
-        auto& l1 = iparams.logits1.value();
-        auto& l2 = iparams.logits2.value();
-        auto& p1 = iparams.probs1.value();
-        auto& p2 = iparams.probs2.value();
         auto& t1 = iparams.tokens1.value();
         auto& t2 = iparams.tokens2.value();
-        assert(l1.size() == t1.size() && l1.size() == p1.size());
-        assert(l2.size() == t2.size() && l2.size() == p2.size());
+        auto& l1 = iparams.logits1.value();
+        auto& l2 = iparams.logits2.value();
+        assert(l2.size() == t2.size());
+        assert(l1.size() == t1.size());
 
         ac::llama::TokenDataVector data1;
-        data1.resize(t1.size());
+        data1.reserve(t1.size());
         for (size_t i = 0; i < t1.size(); i++) {
-            data1[i] = ac::llama::TokenData{
-                .token = t1[i],
-                .logit = l1[i],
-                .prob = p1[i]
-            };
+            data1.emplace_back(ac::llama::TokenData{ t1[i], l1[i] });
         }
 
         ac::llama::TokenDataVector data2;
-        data2.resize(t2.size());
+        data2.reserve(t2.size());
         for (size_t i = 0; i < t2.size(); i++) {
-            data2[i] = ac::llama::TokenData{
-                .token = t2[i],
-                .logit = l2[i],
-                .prob = p2[i]
-            };
+            data2.emplace_back(ac::llama::TokenData{ t2[i], l2[i] });
         }
 
         co_await io.push(Frame_from(sc::StateGeneralInstance::OpCompareTokenData{}, {

ac-local-plugin/example/ep-run.cpp

@@ -78,15 +78,13 @@ int main() try {
 
     auto result2 = llama.call<schema::StateGeneralInstance::OpGetTokenData>({});
 
-    std::cout << "Token Data [0]: " << result2.tokens.value()[0] << ", " << result2.logits.value()[0] << ", " << result2.probs.value()[0] << std::endl;
+    std::cout << "Token Data [0]: " << result2.tokens.value()[0] << ", " << result2.logits.value()[0] << std::endl;
 
     auto result3 = llama.call<schema::StateGeneralInstance::OpCompareTokenData>({
         .tokens1 = result2.tokens,
         .logits1 = result2.logits,
-        .probs1 = result2.probs,
         .tokens2 = result2.tokens,
-        .logits2 = result2.logits,
-        .probs2 = result2.probs
+        .logits2 = result2.logits
     });
 
     std::cout << "Token Data Compare: " << result3.equal.value() << std::endl;

ac-local-plugin/schema/ac/schema/LlamaCpp.hpp

@@ -168,13 +168,11 @@ struct StateGeneralInstance {
         struct Return {
             Field<std::vector<int32_t>> tokens;
             Field<std::vector<float>> logits;
-            Field<std::vector<float>> probs;
 
             template <typename Visitor>
             void visitFields(Visitor& v) {
                 v(tokens, "tokens", "Tokens in the context");
                 v(logits, "logits", "Logits for the tokens");
-                v(probs, "probs", "Probabilities for the tokens");
             }
         };
 
@@ -188,19 +186,15 @@ struct StateGeneralInstance {
         struct Params {
             Field<std::vector<int32_t>> tokens1;
             Field<std::vector<float>> logits1;
-            Field<std::vector<float>> probs1;
             Field<std::vector<int32_t>> tokens2;
             Field<std::vector<float>> logits2;
-            Field<std::vector<float>> probs2;
 
             template <typename Visitor>
             void visitFields(Visitor& v) {
                 v(tokens1, "tokens1", "Tokens in the first set");
                 v(logits1, "logits1", "Logits for the first set");
-                v(probs1, "probs1", "Probabilities for the first set");
                 v(tokens2, "tokens2", "Tokens in the second set");
                 v(logits2, "logits2", "Logits for the second set");
-                v(probs2, "probs2", "Probabilities for the second set");
             }
         };
 

code/ac/llama/LogitComparer.cpp

@@ -4,9 +4,31 @@
 #include "LogitComparer.hpp"
 #include <cmath>
 #include <cassert>
+namespace {
+std::unordered_map<int32_t, float> softmax(const ac::llama::TokenDataVector& data) {
+    std::unordered_map<int32_t, float> result(data.size());
+
+    // Step 1: Find max logit to subtract for numerical stability
+    float maxLogit = data[0].logit;
+
+    // Step 2: Compute exp(logit - maxLogit) for each element
+    float sumExp = 0.0f;
+    for (size_t i = 0; i < data.size(); ++i) {
+        float p = std::exp(data[i].logit - maxLogit);
+        result[data[i].token] = p;
+        sumExp += p;
+    }
 
-namespace ac::llama {
+    // Step 3: Normalize to get probabilities
+    for (auto& val : result) {
+        val.second /= sumExp;
+    }
+
+    return result;
+}
+}
 
+namespace ac::llama {
 // We apply 3 step comparison
 // 1. Compare the euclidean distance of the logits
 //  - If the distance is less than 2% of the max distance, we consider them equal
@@ -25,10 +47,8 @@ bool LogitComparer::compare(const TokenDataVector& data1, const TokenDataVector&
         return false;
     }
 
-    std::unordered_map<int32_t, float> prob_map, prob_map2;
-
-    for (const auto& p : data1) prob_map[p.token] = p.prob;
-    for (const auto& p : data2) prob_map2[p.token] = p.prob;
+    auto prob_map = softmax(data1);
+    auto prob_map2 = softmax(data2);
 
     // Check if at least 80% of the tokens are the same
     float matchingTokens = 0;
@@ -56,7 +76,7 @@ float LogitComparer::logitSimilarity(const TokenDataVector& data1, const TokenDa
     float weightedSimSum = 0.0f;
     float totalWeight = 0.0f;
     for (auto& t : data1) {
-        float weight = t.prob;
+        float weight = t.logit;
         float sim = 0.0f;
         if (l2_map.count(t.token)) {
             sim = 1 - (std::abs(t.logit - l2_map[t.token]) / std::max(t.logit, l2_map[t.token]));

code/ac/llama/Session.cpp

@@ -29,25 +29,7 @@ void fillLogits(TokenDataVector& out, llama_context* lctx) {
     out.resize(vocabSize);
 
     for (llama_token id = 0; id < vocabSize; id++) {
-        out[id] = {id, logits[id], 0.0f};
-    }
-}
-
-static void applySoftMax(TokenDataVector& data) {
-    // Apply softmax to the logits
-    // The vector should be sorted in descending order
-
-    float max_l = data[0].logit;
-    float cum_sum = 0.0f;
-
-    for (size_t i = 0; i < data.size(); ++i) {
-        float p = expf(data[i].logit - max_l);
-        data[i].prob = p;
-        cum_sum += p;
-    }
-
-    for (size_t i = 0; i < data.size(); ++i) {
-        data[i].prob /= cum_sum;
+        out[id] = {id, logits[id]};
     }
 }
 }
@@ -209,9 +191,7 @@ TokenDataVector Session::getSampledTokenData(int32_t topK, float /*topP*/) {
     });
 
     TokenDataVector result;
-    result.insert(result.end(), tempData.begin(), tempData.begin() + topK);
-
-    applySoftMax(result);
+    result.insert(result.end(), tempData.begin(), tempData.begin() + topK);;
 
     return result;
 }

code/ac/llama/Token.hpp

@@ -12,7 +12,6 @@ inline constexpr Token Token_Invalid = -1;
 struct TokenData {
     Token token;
     float logit;
-    float prob;
 };
 
 using TokenDataVector = std::vector<TokenData>;