Add regex capability for tensor selection

Author: EAddario

examples/quantize/quantize.cpp

@@ -319,7 +319,15 @@ static bool string_parse_tensor_type(const char * data, std::vector<tensor_quant
     sep++;
     const std::string qt(sep);
 
-    if (find(ALLOWED_TENSOR_TYPE.begin(), ALLOWED_TENSOR_TYPE.end(), tn) == ALLOWED_TENSOR_TYPE.end()) {
+    bool found = false;
+    for (const auto & allowed : ALLOWED_TENSOR_TYPE) {
+        // check if an allowed tensor exists and it's at the end of the kv string
+        if (tn.length() - allowed.length() == tn.find(allowed)) {
+            found = true;
+            break;
+        }
+    }
+    if (!found) {
         printf("\n%s: invalid tensor name '%s'\n\n", __func__, tn.c_str());
         return false;
     }

src/llama-quant.cpp

@@ -10,6 +10,7 @@
 #include <cinttypes>
 #include <fstream>
 #include <mutex>
+#include <regex>
 #include <thread>
 #include <unordered_map>
 
@@ -795,9 +796,10 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std::
                 // unless the user specifies a type
                 if (params->tensor_types) {
                     const std::vector<tensor_quantization> & tensor_types = *static_cast<const std::vector<tensor_quantization> *>(params->tensor_types);
-                    for (const auto & [name, quant] : tensor_types) {
-                        if (std::string str(tensor->name); str.find(name) != std::string::npos) {
-                            new_type = quant;
+                    for (const auto & [tname, qtype] : tensor_types) {
+                        if (std::regex pattern(tname); std::regex_search(tensor->name, pattern)) {
+                            LLAMA_LOG_DEBUG("(overriding %s -> %s), ", ggml_type_name(new_type), ggml_type_name(qtype));
+                            new_type = qtype;
                             break;
                         }
                     }