implement our version of argmax

inference_test.py

@@ -6,11 +6,27 @@ import jaconv
 from transformers import AutoTokenizer
 import os
 from pathlib import Path
-from typing import Optional, Union
+from typing import Optional, Union, List
+
 
 import re
 import numpy as np
 
+
+# we we're kinda forced to implement the argmax by hand because of two factors:
+# 1- the argmax from this model always takes the last axis, and this is not possible in numcpp , it only goes right or left.
+# 2- numpy's implementation of argmax was too obscure , so we couldn't replicate it in c++ .
+# well, now we can rewrite in c++, and also we dont need to give numpy any credits for this.
+def argmax_last_axis(mat: List[List[int]]) -> List[int]:
+    """Return first max index for each row in matrices of ANY size/shape."""    
+    return [
+        max(
+            ((val, idx) for idx, val in enumerate(row)),
+            default=(-1, -1)
+        )[1]
+        for row in mat
+    ]
+
 def get_most_recent_file(
     directory: Union[str, Path],
     pattern: str = "*",
@@ -60,8 +76,18 @@ class TextProcessor:
     def _remove_specials(self, text: str) -> str:
         """Remove special tokens and consecutive duplicates"""
         # Filter out special tokens
-        return self.special_token_pattern.sub('', text)
+        cleaned = self.special_token_pattern.sub('', text)
         
+        # Remove consecutive duplicates
+        if not cleaned:
+            return ''
+            
+        result = [cleaned[0]]
+        for char in cleaned[1:]:
+            if char != result[-1]:
+                result.append(char)
+        return ''.join(result)
+
 
     def decode(self, indices) -> str:
         """
@@ -80,6 +106,8 @@ class TextProcessor:
 
 
 
+
+
 def preprocess_image(image_path):
     img = cv2.imread(image_path)
     mat_in = ncnn.Mat.from_pixels_resize(
@@ -92,7 +120,7 @@ def preprocess_image(image_path):
 
 def run_inference(mat_in, model_path="model"):
     net = ncnn.Net()
-    net.opt.num_threads = 16
+    net.opt.num_threads = 32
     net.opt.use_vulkan_compute = False
     net.load_param(f"{model_path}/ncnn_model.param")
     net.load_model(f"{model_path}/ncnn_model.bin")
@@ -105,9 +133,10 @@ def run_inference(mat_in, model_path="model"):
         ret, out0 = ex.extract("out0")
         print("Output shape:", out0.shape)
         predictions = np.array(out0)
-        indexed = np.argmax(predictions, axis=-1)
-        print(indexed)
-        return indexed
+      #  ids = np.argmax(predictions, axis=-1)
+        ids = argmax_last_axis(predictions) 
+        print("argmax'ed: ", ids)
+        return ids
     finally:
         del ex, net