Changes to enable rust bindings

.gitignore

@@ -1,8 +1,11 @@
 *.bin
 *.o
 *.npy
+*.a
+*.so
 preprocessor
 nemotron-speech
+nemotron-asr.cpp
 
 # Test binaries
 test_*

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "ggml"]
+	path = ggml
+	url = https://github.com/ggml-org/ggml

Makefile

@@ -1,56 +1,55 @@
 # Makefile for ggml-based NeMo ASR implementation
 
-GGML_DIR = ggml
-GGML_BUILD = $(GGML_DIR)/build
+GGML_DIR ?= ggml
+GGML_BUILD ?= $(GGML_DIR)/build
 
-CXX = g++
-CXXFLAGS = -g -std=c++17 -Wall -Wextra -O2
+CXX ?= g++
+CXXFLAGS ?= -g -std=c++17 -Wall -Wextra -O2
 CXXFLAGS += -I $(GGML_DIR)/include
 CXXFLAGS += -I include
 
-# Check if CUDA backend is available
-CUDA_LIB = $(GGML_BUILD)/src/ggml-cuda/libggml-cuda.so
-CUDA_AVAILABLE = $(shell test -f $(CUDA_LIB) && echo 1 || echo 0)
-
-METAL_LIB = $(GGML_BUILD)/src/ggml-metal/libggml-metal.dylib
-METAL_AVAILABLE = $(shell test -f $(METAL_LIB) && echo 1 || echo 0)
-
-LDFLAGS = -L $(GGML_BUILD)/src
-LDFLAGS += -lggml -lggml-base -lggml-cpu
-LDFLAGS += -Wl,-rpath,$(GGML_BUILD)/src
-LDFLAGS += -lm -lpthread
-
-# Add CUDA support if available
-ifeq ($(CUDA_AVAILABLE),1)
-    CXXFLAGS += -DGGML_USE_CUDA
-    LDFLAGS += -L $(GGML_BUILD)/src/ggml-cuda -lggml-cuda
-    LDFLAGS += -Wl,-rpath,$(GGML_BUILD)/src/ggml-cuda
-    LDFLAGS += -L /usr/local/cuda/lib64 -lcudart -lcublas
-    LDFLAGS += -Wl,-rpath,/usr/local/cuda/lib64
-endif
-
-# Add Metal support if available
-ifeq ($(METAL_AVAILABLE),1)
-	CXXFLAGS += -DGGML_USE_METAL
-	LDFLAGS += -L $(GGML_BUILD)/src/ggml-metal -lggml-metal
-	LDFLAGS += -Wl,-rpath,$(GGML_BUILD)/src/ggml-metal
-	LDFLAGS += -framework Metal -framework Foundation
+# Detect static libraries in GGML build directory
+GGML_STATIC_LIBS := $(wildcard $(GGML_BUILD)/src/*.a)
+
+# Configure linking based on whether static libraries are available
+ifneq ($(GGML_STATIC_LIBS),)
+    # Static linking - use .a files directly
+    LDFLAGS += $(GGML_STATIC_LIBS)
+else
+    # Dynamic linking - use -l flags and rpath
+    LDFLAGS += -L $(GGML_BUILD)/src
+    LDFLAGS += -lggml -lggml-base
+    LDFLAGS += -Wl,-rpath,$(GGML_BUILD)/src
+    LDFLAGS += -Wl,-rpath,$(GGML_BUILD)/bin
 endif
 
 # Source files
 GGML_SRCS = src/nemo-ggml.cpp src/preprocessor.cpp
 GGML_STREAM_SRCS = src/nemo-stream.cpp
+C_WRAPPER_SRCS = src/nemotron_asr_c.cpp
 
 # Original implementation (for comparison tests)
 ORIG_SRCS = src/reference/ggml_weights.cpp src/reference/ops.cpp src/reference/conv_subsampling.cpp src/reference/conformer_modules.cpp src/reference/conformer_encoder.cpp src/reference/rnnt_decoder.cpp src/reference/rnnt_joint.cpp src/reference/greedy_decode.cpp src/reference/tokenizer.cpp
 
-.PHONY: all clean clean_bin test transcribe streaming
+.PHONY: all clean clean_bin test transcribe streaming lib
 
 all: nemotron-asr.cpp
 # test_ggml_weights test_ggml_compute transcribe streaming
 
 streaming: test_streaming nemotron-asr.cpp
 
+# C wrapper library for FFI (static and shared)
+lib: libnemotron_asr.a libnemotron_asr.so
+
+libnemotron_asr.a: $(GGML_SRCS:.cpp=.o) $(GGML_STREAM_SRCS:.cpp=.o) $(C_WRAPPER_SRCS:.cpp=.o)
+	ar rcs $@ $^
+
+libnemotron_asr.so: $(GGML_SRCS) $(GGML_STREAM_SRCS) $(C_WRAPPER_SRCS)
+	$(CXX) $(CXXFLAGS) -fPIC -shared $^ $(LDFLAGS) -o $@
+
+%.o: %.cpp
+	$(CXX) $(CXXFLAGS) -fPIC -c $< -o $@
+
 # Test weight loading
 test_ggml_weights: tests/test_weights.cpp $(GGML_SRCS) $(ORIG_SRCS)
 	$(CXX) $(CXXFLAGS) $^ $(LDFLAGS) -o $@
@@ -85,6 +84,8 @@ nemotron-asr.cpp: src/transcribe_stream.cpp $(GGML_SRCS) $(GGML_STREAM_SRCS)
 
 clean:
 	rm -f test_ggml_weights test_ggml_compute precompute_encoder_ref transcribe test_streaming transcribe_stream test_python_ref test_preprocessor
+	rm -f libnemotron_asr.a libnemotron_asr.so
+	rm -f src/*.o
 
 clean_bin:
 	rm my_bin/*

README.md

@@ -36,10 +36,9 @@ Additionally, those changes also help with quantization, which has requirements
 
 ## Development
 
-To develop this project you need to clone [ggml-org/ggml](https://github.com/ggml-org/ggml) to this directory, then build ggml with all your favorite settings. TODO: Maybe make this process more friendly? PRs welcome.
-
+First, build ggml with all your favorite settings.
 ```bash
-git clone https://github.com/ggml-org/ggml.git
+git subdmodule update --init
 mkdir ggml/build
 cd ggml/build
 cmake ..
@@ -52,12 +51,14 @@ Once you have built ggml, you can build this binary
 make nemotron-asr.cpp
 ```
 
+All of the ggml backends that you built with should be available to `nemotron-asr.cpp` through the `--backend` flag.
+
 ## Comparison to [whisper.cpp](https://github.com/ggml-org/whisper.cpp)?
 
 1) Whisper operates on 30s audio chunks, so it is not feasible to use in interactive applications.
 NeMotron's ASR model has configurable latency (from 80ms to 1.12s), which can traded for quality and speed. (bigger lookahead gives better quality and bigger chunks require less work)
 2) Whisper has troubles on long audio streams, where NeMotron's ASR is a streaming model, works for infinite streams. It does get stuck in lowercase mode though.
-3) Quiality seems to be better
+3) Quality seems to be better
 
 ## License
 

src/nemo-ggml.cpp

@@ -31,60 +31,60 @@ static void compute_pos_emb(float * data, int max_len, int d_model) {
     }
 }
 
-// Initialize backend based on requested type
-static bool init_backend(nemo_model & model, nemo_backend_type backend_type) {
+// Initialize backend based on requested name (using dynamic backend loading)
+// backend_name: specific backend to use (e.g., "CPU", "CUDA", "Vulkan", "Metal")
+// If nullptr, uses the first available backend (highest priority)
+static bool init_backend(nemo_model & model, const char * backend_name) {
     model.backend = nullptr;
-    model.backend_type = NEMO_BACKEND_CPU;  // default
+    model.backend_name = "";
 
-#ifdef GGML_USE_METAL
-    if (backend_type == NEMO_BACKEND_METAL || backend_type == NEMO_BACKEND_AUTO) {
-        model.backend = ggml_backend_metal_init();
+    // Load all available backend modules
+    ggml_backend_load_all();
+
+    if (backend_name != nullptr && backend_name[0] != '\0') {
+        // User specified a backend
+        model.backend = ggml_backend_init_by_name(backend_name, nullptr);
         if (!model.backend) {
-            fprintf(stderr, "Failed to initialize Metal backend\n");
-            assert(false);
+            fprintf(stderr, "%s: failed to initialize %s backend\n", __func__, backend_name);
+            return false;
+        }
+        model.backend_name = backend_name;
+        printf("%s: using %s backend\n", __func__, backend_name);
+    } else {
+        // Auto-select: use first available backend (backends are registered in priority order)
+        size_t dev_count = ggml_backend_dev_count();
+        if (dev_count == 0) {
+            fprintf(stderr, "%s: no backends available\n", __func__);
+            return false;
         }
-    }
-#endif
 
-#ifdef GGML_USE_CUDA
-    if (backend_type == NEMO_BACKEND_CUDA || backend_type == NEMO_BACKEND_AUTO) {
-        int n_devices = ggml_backend_cuda_get_device_count();
-        if (n_devices > 0) {
-            model.backend = ggml_backend_cuda_init(0);  // use first CUDA device
+        // Try devices in order until one initializes successfully
+        for (size_t i = 0; i < dev_count; i++) {
+            ggml_backend_dev_t dev = ggml_backend_dev_get(i);
+            const char * dev_name = ggml_backend_dev_name(dev);
+
+            model.backend = ggml_backend_dev_init(dev, nullptr);
             if (model.backend) {
-                model.backend_type = NEMO_BACKEND_CUDA;
-                char desc[256];
-                ggml_backend_cuda_get_device_description(0, desc, sizeof(desc));
-                printf("%s: using CUDA backend (%s)\n", __func__, desc);
-
-                size_t free_mem, total_mem;
-                ggml_backend_cuda_get_device_memory(0, &free_mem, &total_mem);
-                printf("%s: CUDA memory: %.1f / %.1f GB available\n", __func__,
-                       free_mem / 1e9, total_mem / 1e9);
+                model.backend_name = dev_name;
+                printf("%s: using %s backend (auto-selected)\n", __func__, dev_name);
+                break;
             }
         }
-    }
-#endif
 
-    // Fall back to CPU if CUDA not available or not requested
-    if (!model.backend) {
-        if (backend_type == NEMO_BACKEND_CUDA) {
-            fprintf(stderr, "%s: CUDA backend requested but not available\n", __func__);
+        if (!model.backend) {
+            fprintf(stderr, "%s: failed to initialize any backend\n", __func__);
             return false;
         }
-        model.backend = ggml_backend_cpu_init();
-        model.backend_type = NEMO_BACKEND_CPU;
-        printf("%s: using CPU backend\n", __func__);
     }
 
-    return model.backend != nullptr;
+    return true;
 }
 
-bool nemo_model_load(const std::string & path, nemo_model & model, nemo_backend_type backend_type) {
+bool nemo_model_load(const std::string & path, nemo_model & model, const char * backend_name) {
     // printf("%s: loading model from '%s'\n", __func__, path.c_str());
 
     // Initialize backend
-    if (!init_backend(model, backend_type)) {
+    if (!init_backend(model, backend_name)) {
         fprintf(stderr, "%s: failed to initialize backend\n", __func__);
         return false;
     }
@@ -388,13 +388,13 @@ bool nemo_model_load(const std::string & path, nemo_model & model, nemo_backend_
 }
 
 struct nemo_context * nemo_init(const char * model_path) {
-    return nemo_init_with_backend(model_path, NEMO_BACKEND_AUTO);
+    return nemo_init_with_backend(model_path, nullptr);
 }
 
-struct nemo_context * nemo_init_with_backend(const char * model_path, nemo_backend_type backend) {
+struct nemo_context * nemo_init_with_backend(const char * model_path, const char * backend_name) {
     struct nemo_context * ctx = new nemo_context();
 
-    if (!nemo_model_load(model_path, ctx->model, backend)) {
+    if (!nemo_model_load(model_path, ctx->model, backend_name)) {
         delete ctx;
         return nullptr;
     }
@@ -433,12 +433,8 @@ struct nemo_context * nemo_init_with_backend(const char * model_path, nemo_backe
 }
 
 const char * nemo_get_backend_name(struct nemo_context * ctx) {
-    if (!ctx) return "unknown";
-    switch (ctx->model.backend_type) {
-        case NEMO_BACKEND_CUDA: return "CUDA";
-        case NEMO_BACKEND_CPU:  return "CPU";
-        default: return "unknown";
-    }
+    if (!ctx || ctx->model.backend_name.empty()) return "unknown";
+    return ctx->model.backend_name.c_str();
 }
 
 void nemo_free(struct nemo_context * ctx) {

src/nemo-ggml.h

@@ -7,13 +7,8 @@
 #include "ggml-backend.h"
 #include "gguf.h"
 
-#ifdef GGML_USE_METAL
-#include "ggml-metal.h"
-#endif
-
-#ifdef GGML_USE_CUDA
-#include "ggml-cuda.h"
-#endif
+// Backend headers not needed with dynamic loading
+// Backends are loaded via ggml_backend_init_by_name() or ggml_backend_init_by_type()
 
 #include <cstdint>
 #include <string>
@@ -22,14 +17,6 @@
 
 struct timed_token;
 
-// Backend type for inference
-enum nemo_backend_type {
-    NEMO_BACKEND_CPU = 0,
-    NEMO_BACKEND_CUDA = 1,
-    NEMO_BACKEND_METAL = 2,
-    NEMO_BACKEND_AUTO = 3,  // Auto-detect: prefer CUDA if available
-};
-
 // Forward declaration
 struct nemo_preprocessor;
 
@@ -181,7 +168,7 @@ struct nemo_model {
     struct ggml_context * ctx_w;          // weights context
     ggml_backend_t backend;
     ggml_backend_buffer_t buffer_w;
-    nemo_backend_type backend_type;       // which backend is in use
+    std::string backend_name;             // which backend is in use
 
     // Tensor name mapping for loading
     std::map<std::string, struct ggml_tensor *> tensors;
@@ -227,19 +214,21 @@ struct nemo_context {
 };
 
 // API functions
-// Initialize with automatic backend selection (prefers CUDA if available)
+// Initialize with automatic backend selection (uses first available backend from GGML registry)
 struct nemo_context * nemo_init(const char * model_path);
 
-// Initialize with specific backend
-struct nemo_context * nemo_init_with_backend(const char * model_path, nemo_backend_type backend);
+// Initialize with specific backend by name (e.g., "CPU", "CUDA", "Vulkan", "Metal")
+// Pass nullptr for backend_name to use automatic selection
+struct nemo_context * nemo_init_with_backend(const char * model_path, const char * backend_name);
 
 void nemo_free(struct nemo_context * ctx);
 
 // Get current backend name
 const char * nemo_get_backend_name(struct nemo_context * ctx);
 
 // Load model weights from file (with backend selection)
-bool nemo_model_load(const std::string & path, nemo_model & model, nemo_backend_type backend = NEMO_BACKEND_AUTO);
+// Pass nullptr for backend_name to use automatic selection
+bool nemo_model_load(const std::string & path, nemo_model & model, const char * backend_name = nullptr);
 
 // Build computation graphs
 struct ggml_cgraph * nemo_build_encoder_graph(