Add GetOutputTensorDataRaw method to AxclModel and update FireRedASR CTC model to use it

Author: LittleMouse

sherpa-onnx/csrc/axcl/axcl-model.cc

@@ -208,6 +208,31 @@ class AxclModel::Impl {
     return {};
   }
 
+  std::vector<uint8_t> GetOutputTensorDataRaw(const std::string &name) const {
+    for (size_t i = 0; i < output_tensor_names_.size(); ++i) {
+      if (output_tensor_names_[i] == name) {
+        size_t bytes = output_tensors_[i].Size();
+        std::vector<uint8_t> out(bytes);
+
+        auto ret = axclrtMemcpy(out.data(), output_tensors_[i].Get(), bytes,
+                                AXCL_MEMCPY_DEVICE_TO_HOST);
+        if (ret != 0) {
+          SHERPA_ONNX_LOGE(
+              "Failed to call axclrtMemcpy(). tensor name: '%s', return code: "
+              "%d",
+              name.c_str(), static_cast<int32_t>(ret));
+          return {};
+        }
+
+        return out;
+      }
+    }
+
+    SHERPA_ONNX_LOGE("Found no tensor with name: '%s'", name.c_str());
+
+    return {};
+  }
+
   bool Run() const {
     uint32_t group = 0;
     auto ret =
@@ -434,6 +459,11 @@ std::vector<float> AxclModel::GetOutputTensorData(
   return impl_->GetOutputTensorData(name);
 }
 
+std::vector<uint8_t> AxclModel::GetOutputTensorDataRaw(
+    const std::string &name) const {
+  return impl_->GetOutputTensorDataRaw(name);
+}
+
 bool AxclModel::Run() const { return impl_->Run(); }
 
 bool AxclModel::IsInitialized() const { return impl_->IsInitialized(); }

sherpa-onnx/csrc/axcl/axcl-model.h

@@ -36,6 +36,8 @@ class AxclModel {
 
   std::vector<float> GetOutputTensorData(const std::string &name) const;
 
+  std::vector<uint8_t> GetOutputTensorDataRaw(const std::string &name) const;
+
   bool Run() const;
   bool IsInitialized() const;
 

sherpa-onnx/csrc/axcl/offline-fire-red-asr-ctc-model-axcl.cc

@@ -5,6 +5,8 @@
 #include "sherpa-onnx/csrc/axcl/offline-fire-red-asr-ctc-model-axcl.h"
 
 #include <algorithm>
+#include <cmath>
+#include <cstring>
 #include <memory>
 #include <string>
 #include <utility>
@@ -64,28 +66,47 @@ class OfflineFireRedAsrCtcModelAxcl::Impl {
     int32_t expected_frames = expected_shape[1];
 
     int32_t valid_frames = std::min<int32_t>(num_frames, expected_frames);
+    valid_frames = std::min<int32_t>(valid_frames,
+                                     static_cast<int32_t>(p_features_length[0]));
     std::vector<float> padded_features(expected_frames * feat_dim, 0.0f);
     std::copy(p_features, p_features + valid_frames * feat_dim,
               padded_features.begin());
 
     std::vector<int32_t> speech_length = {valid_frames};
 
-    model_->SetInputTensorData(model_->InputTensorNames()[0],
-                               padded_features.data(), padded_features.size());
-    model_->SetInputTensorData(model_->InputTensorNames()[1],
-                               speech_length.data(), speech_length.size());
+    if (!model_->SetInputTensorData(model_->InputTensorNames()[0],
+                                    padded_features.data(),
+                                    padded_features.size()) ||
+        !model_->SetInputTensorData(model_->InputTensorNames()[1],
+                                    speech_length.data(),
+                                    speech_length.size())) {
+      SHERPA_ONNX_LOGE("Failed to set input tensors for axcl FireRedASR CTC");
+      SHERPA_ONNX_EXIT(-1);
+    }
 
-    model_->Run();
+    if (!model_->Run()) {
+      SHERPA_ONNX_LOGE("Failed to run axcl FireRedASR CTC model");
+      SHERPA_ONNX_EXIT(-1);
+    }
 
     auto out_logits =
         model_->GetOutputTensorData(model_->OutputTensorNames()[0]);
-    auto out_lengths =
-        model_->GetOutputTensorData(model_->OutputTensorNames()[1]);
+    auto out_lengths_raw =
+        model_->GetOutputTensorDataRaw(model_->OutputTensorNames()[1]);
 
     auto out_shape = model_->TensorShape(model_->OutputTensorNames()[0]);
     int32_t out_frames = out_shape[1];
     int32_t vocab_size = out_shape[2];
 
+    if (static_cast<int32_t>(out_logits.size()) != out_frames * vocab_size) {
+      SHERPA_ONNX_LOGE(
+          "Unexpected logits size from axcl FireRedASR CTC. Got %d values. "
+          "Expected %d x %d = %d",
+          static_cast<int32_t>(out_logits.size()), out_frames, vocab_size,
+          out_frames * vocab_size);
+      SHERPA_ONNX_EXIT(-1);
+    }
+
     std::array<int64_t, 3> logits_shape = {1, out_frames, vocab_size};
     Ort::Value logits = Ort::Value::CreateTensor<float>(
         allocator_, logits_shape.data(), logits_shape.size());
@@ -96,7 +117,7 @@ class OfflineFireRedAsrCtcModelAxcl::Impl {
     Ort::Value lengths = Ort::Value::CreateTensor<int64_t>(
         allocator_, lengths_shape.data(), lengths_shape.size());
     int64_t *p_lengths = lengths.GetTensorMutableData<int64_t>();
-    p_lengths[0] = static_cast<int64_t>(out_lengths[0]);
+    p_lengths[0] = ReadOutputLength(out_lengths_raw, out_frames);
 
     std::vector<Ort::Value> ans;
     ans.push_back(std::move(logits));
@@ -131,6 +152,55 @@ class OfflineFireRedAsrCtcModelAxcl::Impl {
   }
 
  private:
+  static int64_t ReadOutputLength(const std::vector<uint8_t> &raw,
+                                  int32_t max_frames) {
+    if (raw.empty()) {
+      SHERPA_ONNX_LOGE(
+          "The output-length tensor from axcl FireRedASR CTC is empty");
+      SHERPA_ONNX_EXIT(-1);
+    }
+
+    auto in_range = [max_frames](int64_t value) {
+      return 0 <= value && value <= max_frames;
+    };
+
+    if (raw.size() == sizeof(int64_t)) {
+      int64_t value = 0;
+      std::memcpy(&value, raw.data(), sizeof(value));
+      if (in_range(value)) {
+        return value;
+      }
+    }
+
+    if (raw.size() >= sizeof(int32_t)) {
+      int32_t value_i32 = 0;
+      std::memcpy(&value_i32, raw.data(), sizeof(value_i32));
+      if (in_range(value_i32)) {
+        return value_i32;
+      }
+
+      uint32_t value_u32 = 0;
+      std::memcpy(&value_u32, raw.data(), sizeof(value_u32));
+      if (value_u32 <= static_cast<uint32_t>(max_frames)) {
+        return value_u32;
+      }
+
+      float value_f32 = 0;
+      std::memcpy(&value_f32, raw.data(), sizeof(value_f32));
+      if (std::isfinite(value_f32) && value_f32 >= 0 &&
+          value_f32 <= static_cast<float>(max_frames)) {
+        return static_cast<int64_t>(value_f32);
+      }
+    }
+
+    SHERPA_ONNX_LOGE(
+        "Failed to interpret the output-length tensor from axcl FireRedASR "
+        "CTC. Byte size: %d. Max frames: %d",
+        static_cast<int32_t>(raw.size()), max_frames);
+    SHERPA_ONNX_EXIT(-1);
+    return 0;
+  }
+
   void Init() {
     if (!model_->IsInitialized()) {
       SHERPA_ONNX_LOGE("Failed to initialize the model with '%s'",