Add functionality for LabelScorers implemented in python

src/Python/LabelScorer.cc

@@ -0,0 +1,176 @@
+/** Copyright 2025 RWTH Aachen University. All rights reserved.
+ *
+ *  Licensed under the RWTH ASR License (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *      http://www.hltpr.rwth-aachen.de/rwth-asr/rwth-asr-license.html
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+#include "LabelScorer.hh"
+
+#include <pybind11/gil.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include "ScoringContext.hh"
+
+namespace py = pybind11;
+
+namespace Python {
+
+PythonLabelScorer::PythonLabelScorer(Core::Configuration const& config)
+        : Core::Component(config),
+          Precursor(config) {
+}
+
+void PythonLabelScorer::reset() {
+    PYBIND11_OVERRIDE_PURE(void, LabelScorer, reset);
+}
+
+void PythonLabelScorer::signalNoMoreFeatures() {
+    PYBIND11_OVERRIDE_PURE_NAME(
+            void,
+            LabelScorer,
+            "signal_no_more_features",
+            signalNoMoreFeatures);
+}
+
+Nn::ScoringContextRef PythonLabelScorer::getInitialScoringContext() {
+    py::gil_scoped_acquire gil;
+    // Store `py::object` from virtual python call in a `PythonScoringContext`
+    return Core::ref(new PythonScoringContext(getInitialPythonScoringContext()));
+}
+
+py::object PythonLabelScorer::getInitialPythonScoringContext() {
+    PYBIND11_OVERRIDE_PURE_NAME(
+            py::object,
+            Nn::LabelScorer,
+            "get_initial_scoring_context",
+            getInitialPythonScoringContext);
+}
+
+Nn::ScoringContextRef PythonLabelScorer::extendedScoringContext(Request const& request) {
+    auto*                  pythonScoringContext = dynamic_cast<PythonScoringContext const*>(request.context.get());
+    py::gil_scoped_acquire gil;
+    // Store `py::object` from virtual python call in a `PythonScoringContext`
+    auto newScoringContext = extendedPythonScoringContext(pythonScoringContext->object, request.nextToken, request.transitionType);
+    return Core::ref(new PythonScoringContext(std::move(newScoringContext)));
+}
+
+py::object PythonLabelScorer::extendedPythonScoringContext(py::object const& context, Nn::LabelIndex nextToken, TransitionType transitionType) {
+    PYBIND11_OVERRIDE_PURE_NAME(
+            py::object,
+            Nn::LabelScorer,
+            "extended_scoring_context",
+            extendedPythonScoringContext,
+            context,
+            nextToken,
+            transitionType);
+}
+
+void PythonLabelScorer::addInput(Nn::DataView const& input) {
+    // Call batched version
+    addInputs(input, 1);
+}
+
+void PythonLabelScorer::addInputs(Nn::DataView const& input, size_t nTimesteps) {
+    py::gil_scoped_acquire gil;
+
+    // Convert `input` to a `py::array` for virtual python call
+    ssize_t featureDimSize = input.size() / nTimesteps;
+
+    py::array_t<f32> inputArray(
+            {static_cast<ssize_t>(nTimesteps), featureDimSize},
+            {sizeof(f32) * featureDimSize, sizeof(f32)},
+            input.data());
+
+    addPythonInputs(inputArray);
+}
+
+void PythonLabelScorer::addPythonInputs(py::array const& inputs) {
+    PYBIND11_OVERRIDE_PURE_NAME(
+            void,
+            Nn::LabelScorer,
+            "add_inputs",
+            addPythonInputs,
+            inputs);
+}
+
+std::optional<Nn::LabelScorer::ScoreWithTime> PythonLabelScorer::computeScoreWithTime(Request const& request) {
+    // Extract the underlying `py::object` from ScoringContext in `request` to supply them to the virtual python call
+    auto* pythonScoringContext = dynamic_cast<PythonScoringContext const*>(request.context.get());
+
+    std::vector<py::object>     contexts        = {pythonScoringContext->object};
+    std::vector<Nn::LabelIndex> nextTokens      = {request.nextToken};
+    std::vector<TransitionType> transitionTypes = {request.transitionType};
+
+    py::gil_scoped_acquire gil;
+
+    // Call batched version
+    if (auto result = computePythonScoresWithTimes(contexts, nextTokens, transitionTypes)) {
+        verify(result->size() == 1);
+        ScoreWithTime scoreWithTime{result->front().first, result->front().second};
+        return scoreWithTime;
+    }
+
+    return {};
+}
+
+std::optional<Nn::LabelScorer::ScoresWithTimes> PythonLabelScorer::computeScoresWithTimes(std::vector<Request> const& requests) {
+    std::vector<py::object>     contexts;
+    std::vector<Nn::LabelIndex> nextTokens;
+    std::vector<TransitionType> transitionTypes;
+
+    contexts.reserve(requests.size());
+    nextTokens.reserve(requests.size());
+    transitionTypes.reserve(requests.size());
+
+    // Extract the underlying `py::object`s from ScoringContexts in `requests` to supply them to the virtual python call
+    for (auto const& request : requests) {
+        auto* pythonScoringContext = dynamic_cast<PythonScoringContext const*>(request.context.get());
+        contexts.push_back(pythonScoringContext->object);
+        nextTokens.push_back(request.nextToken);
+        transitionTypes.push_back(request.transitionType);
+    }
+
+    py::gil_scoped_acquire gil;
+
+    if (auto result = computePythonScoresWithTimes(contexts, nextTokens, transitionTypes)) {
+        verify(result->size() == requests.size());
+        ScoresWithTimes scoresWithTimes;
+        scoresWithTimes.scores.reserve(result->size());
+        for (auto const& [score, timeframe] : *result) {
+            scoresWithTimes.scores.push_back(score);
+            scoresWithTimes.timeframes.push_back(timeframe);
+        }
+        return scoresWithTimes;
+    }
+
+    return {};
+}
+
+std::optional<std::vector<std::pair<Nn::LabelScorer::Score, Speech::TimeframeIndex>>> PythonLabelScorer::computePythonScoresWithTimes(std::vector<py::object> const& contexts, std::vector<Nn::LabelIndex> const& nextTokens, std::vector<TransitionType> const& transitionTypes) {
+    using returnType = std::optional<std::vector<std::pair<Nn::LabelScorer::Score, Speech::TimeframeIndex>>>;  // Macro can't handle types with commas inside properly
+    PYBIND11_OVERRIDE_PURE_NAME(
+            returnType,
+            Nn::LabelScorer,
+            "compute_scores_with_times",
+            computePythonScoresWithTimes,
+            contexts,
+            nextTokens,
+            transitionTypes);
+}
+
+void PythonLabelScorer::setInstance(py::object const& instance) {
+    py::gil_scoped_acquire gil;
+    pyInstance_ = instance;
+}
+
+}  // namespace Python

src/Python/LabelScorer.hh

@@ -0,0 +1,80 @@
+/** Copyright 2025 RWTH Aachen University. All rights reserved.
+ *
+ *  Licensed under the RWTH ASR License (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *      http://www.hltpr.rwth-aachen.de/rwth-asr/rwth-asr-license.html
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+#ifndef PYTHON_LABEL_SCORER_HH
+#define PYTHON_LABEL_SCORER_HH
+
+#include <pybind11/pybind11.h>
+
+#include <Nn/LabelScorer/LabelScorer.hh>
+
+namespace py = pybind11;
+
+namespace Python {
+
+/*
+ * Trampoline class that is used in order to expose the LabelScorer class via pybind.
+ * It mainly specifies the signatures of abstract methods that need to be implemented in python
+ * and performs conversion between "C++ types" such as `DataView` and `ScoringContext`
+ * and "Python types" such as `py::array` and `py::object`.
+ *
+ * See https://pybind11.readthedocs.io/en/stable/advanced/classes.html for official documentation
+ * on the "trampoline" pattern.
+ */
+class PythonLabelScorer : public Nn::LabelScorer {
+public:
+    using Precursor = Nn::LabelScorer;
+
+    PythonLabelScorer(Core::Configuration const& config);
+    virtual ~PythonLabelScorer() = default;
+
+    // Must be overridden in python by name "reset"
+    virtual void reset() override;
+
+    // Can be overridden in python. No-op per default.
+    virtual void signalNoMoreFeatures() override;
+
+    // Must be overridden in python by name "get_initial_scoring_context"
+    virtual Nn::ScoringContextRef getInitialScoringContext() override;
+    virtual py::object            getInitialPythonScoringContext();
+
+    // Must be overridden in python by name "extended_scoring_context"
+    virtual Nn::ScoringContextRef extendedScoringContext(Request const& request) override;
+    virtual py::object            extendedPythonScoringContext(py::object const& context, Nn::LabelIndex nextToken, TransitionType transitionType);
+
+    // Calls batched version with `nTimesteps = 1`
+    virtual void addInput(Nn::DataView const& input) override;
+
+    // Must be overridden in python by name "add_inputs"
+    virtual void addInputs(Nn::DataView const& input, size_t nTimesteps) override;
+    virtual void addPythonInputs(py::array const& inputs);
+
+    // Calls batched version
+    virtual std::optional<ScoreWithTime> computeScoreWithTime(Request const& request) override;
+
+    // Must be overridden in python by name "compute_scores_with_times"
+    virtual std::optional<ScoresWithTimes>                                       computeScoresWithTimes(std::vector<Request> const& requests) override;
+    virtual std::optional<std::vector<std::pair<Score, Speech::TimeframeIndex>>> computePythonScoresWithTimes(std::vector<py::object> const& contexts, std::vector<Nn::LabelIndex> const& nextTokens, std::vector<TransitionType> const& transitionTypes);
+
+    // Keep track of python object as a member to make sure it doesn't get garbage collected
+    void setInstance(py::object const& instance);
+
+protected:
+    py::object pyInstance_;  // Hold the Python wrapper
+};
+
+}  // namespace Python
+
+#endif  // PYTHON_LABEL_SCORER_HH

src/Python/Makefile

@@ -20,7 +20,10 @@ LIBPYTHON_O    = \
 		$(OBJDIR)/AllophoneStateFsaBuilder.o \
 		$(OBJDIR)/Configuration.o            \
 		$(OBJDIR)/Init.o                     \
+		$(OBJDIR)/LabelScorer.o              \
 		$(OBJDIR)/Numpy.o                    \
+		$(OBJDIR)/ScoringContext.o           \
+		$(OBJDIR)/Search.o                   \
 		$(OBJDIR)/Utilities.o
 
 CHECK_O		= $(OBJDIR)/check.o 			\

src/Python/ScoringContext.cc

@@ -0,0 +1,31 @@
+/** Copyright 2025 RWTH Aachen University. All rights reserved.
+ *
+ *  Licensed under the RWTH ASR License (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *      http://www.hltpr.rwth-aachen.de/rwth-asr/rwth-asr-license.html
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+#include "ScoringContext.hh"
+
+namespace Python {
+
+size_t PythonScoringContext::hash() const {
+    return py::hash(py::cast<py::handle>(object));
+}
+
+bool PythonScoringContext::isEqual(Nn::ScoringContextRef const& other) const {
+    auto* otherPtr = dynamic_cast<const PythonScoringContext*>(other.get());
+
+    py::gil_scoped_acquire gil;
+    return object.equal(py::cast<py::handle>(otherPtr->object));
+}
+
+}  // namespace Python

src/Python/ScoringContext.hh

@@ -0,0 +1,47 @@
+/** Copyright 2025 RWTH Aachen University. All rights reserved.
+ *
+ *  Licensed under the RWTH ASR License (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *      http://www.hltpr.rwth-aachen.de/rwth-asr/rwth-asr-license.html
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+#ifndef PYTHON_SCORING_CONTEXT_HH
+#define PYTHON_SCORING_CONTEXT_HH
+
+#include <pybind11/pybind11.h>
+
+#include <Nn/LabelScorer/ScoringContext.hh>
+
+namespace py = pybind11;
+
+namespace Python {
+
+/*
+ * Scoring context containing some arbitrary (hashable) python object
+ */
+struct PythonScoringContext : public Nn::ScoringContext {
+    py::object object;
+
+    PythonScoringContext()
+            : object(py::none()) {}
+
+    PythonScoringContext(py::object&& object)
+            : object(object) {}
+
+    bool   isEqual(Nn::ScoringContextRef const& other) const;
+    size_t hash() const;
+};
+
+typedef Core::Ref<const PythonScoringContext> PythonScoringContextRef;
+
+}  // namespace Python
+
+#endif  // PYTHON_SCORING_CONTEXT_HH