From 38055fae657aaa6fe0d48b18e36ce4167826011e Mon Sep 17 00:00:00 2001
From: Aleksandras Kostarevas <aleks076@protonmail.com>
Date: Wed, 13 Mar 2024 16:11:52 -0500
Subject: [PATCH] Add other workaround

---
 ...futo_inputmethod_latin_xlm_LanguageModel.cpp | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/native/jni/org_futo_inputmethod_latin_xlm_LanguageModel.cpp b/native/jni/org_futo_inputmethod_latin_xlm_LanguageModel.cpp
index fddc37105..11b15e71f 100644
--- a/native/jni/org_futo_inputmethod_latin_xlm_LanguageModel.cpp
+++ b/native/jni/org_futo_inputmethod_latin_xlm_LanguageModel.cpp
@@ -556,6 +556,7 @@ struct LanguageModelState {
             return { };
         }
 
+        // TODO: This should really not be here
         is_bugged = is_bugged && logits[561] < -990.0f && logits[561] > -1100.0f;
         if(is_bugged) {
             AKLOGE("Detected bug!!!! Trying to mitigate. Let's just reset cache and exit");
@@ -592,6 +593,22 @@ struct LanguageModelState {
             );
         }
 
+        // TODO: This should really not be here
+        is_bugged = true;
+        for(const auto &seq : sequences) {
+            if(seq.second.tokens.front() > 48 || seq.first != sequences[0].first) {
+                is_bugged = false;
+                break;
+            }
+        }
+        if(is_bugged) {
+            AKLOGE("Detected bug2!!!! Trying to mitigate. Let's just reset cache and exit");
+            llama_kv_cache_seq_rm(ctx, -1, -1, -1);
+            model->transformerContext.active_context = { };
+            return { };
+        }
+
+
         for (auto &sequence: sequences) {
             if (sequence.second.seq_id == 0) continue;