rlx_omnicoder/
lib.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! OmniCoder family runner.
17//!
18//! Tesslate/OmniCoder-9B and siblings are Qwen3-coder-shaped and ship
19//! as `general.architecture = qwen3` in their GGUF converters. This
20//! crate is a thin wrapper over [`rlx_qwen3::Qwen3Runner`] with arch
21//! validation so users get a discoverable per-family entry point and
22//! misrouted weights fail loudly at `build()`.
23
24use anyhow::{Context, Result, bail};
25use rlx_llama_base::LlamaBaseConfig;
26use std::path::{Path, PathBuf};
27
28pub use rlx_qwen3::{Qwen3ConfigSource, Qwen3Runner, Qwen3RunnerBuilder};
29
30pub const PLAN_MILESTONE: &str = "M4";
31pub const FAMILY: &str = "OmniCoder";
32
33pub struct OmniCoderRunner {
34    inner: Qwen3Runner,
35    config: LlamaBaseConfig,
36}
37
38impl OmniCoderRunner {
39    pub fn builder() -> OmniCoderRunnerBuilder {
40        OmniCoderRunnerBuilder::default()
41    }
42    pub fn config(&self) -> &LlamaBaseConfig {
43        &self.config
44    }
45    pub fn inner(&self) -> &Qwen3Runner {
46        &self.inner
47    }
48    pub fn inner_mut(&mut self) -> &mut Qwen3Runner {
49        &mut self.inner
50    }
51    pub fn generate_packed(
52        &mut self,
53        prompt_ids: &[u32],
54        n_new: usize,
55        on_token: impl FnMut(u32),
56    ) -> Result<Vec<u32>> {
57        self.inner.generate_packed(prompt_ids, n_new, on_token)
58    }
59}
60
61#[derive(Debug, Clone, Default)]
62pub struct OmniCoderRunnerBuilder {
63    weights: Option<PathBuf>,
64    inner: Qwen3RunnerBuilder,
65}
66
67impl OmniCoderRunnerBuilder {
68    pub fn weights(mut self, path: impl Into<PathBuf>) -> Self {
69        let p: PathBuf = path.into();
70        self.weights = Some(p.clone());
71        self.inner = self.inner.weights(p);
72        self
73    }
74    pub fn max_seq(mut self, n: usize) -> Self {
75        self.inner = self.inner.max_seq(n);
76        self
77    }
78    pub fn packed_weights(mut self, on: bool) -> Self {
79        self.inner = self.inner.packed_weights(on);
80        self
81    }
82    pub fn build(self) -> Result<OmniCoderRunner> {
83        let weights = self
84            .weights
85            .as_ref()
86            .ok_or_else(|| anyhow::anyhow!("weights path required"))?
87            .clone();
88        let config = LlamaBaseConfig::from_gguf_path(&weights)
89            .with_context(|| format!("rlx-omnicoder: parse {weights:?}"))?;
90        if !matches!(config.arch.as_str(), "qwen3" | "qwen2") {
91            bail!(
92                "rlx-omnicoder: expected `general.architecture = qwen3` (OmniCoder uses \
93                 Qwen3 arch); got `{}` at {weights:?}",
94                config.arch
95            );
96        }
97        let inner = self
98            .inner
99            .build()
100            .context("rlx-omnicoder: building underlying Qwen3Runner")?;
101        Ok(OmniCoderRunner { inner, config })
102    }
103}
104
105pub fn cli_run(args: &[String]) -> Result<()> {
106    if let Some(first) = args.iter().position(|a| a == "--weights") {
107        if let Some(path) = args.get(first + 1) {
108            let cfg = LlamaBaseConfig::from_gguf_path(Path::new(path))
109                .with_context(|| format!("rlx-omnicoder: parse {path}"))?;
110            if !matches!(cfg.arch.as_str(), "qwen3" | "qwen2") {
111                bail!(
112                    "rlx-omnicoder: {path}: GGUF arch = `{}`, expected `qwen3` (OmniCoder)",
113                    cfg.arch
114                );
115            }
116        }
117    }
118    rlx_qwen3::cli::run(args)
119}
rlx_omnicoder/lib.rs

rlx_omnicoder/
lib.rs
