Skip to main content

rlx_gpt_oss/
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//! gpt-oss-20b runner.
17//!
18//! gpt-oss ships as `general.architecture = gpt-oss` (or `gpt_oss`)
19//! in its GGUF converters. The 20B variant is an attention-only
20//! Llama-shaped decoder; this crate is a thin wrapper over
21//! [`rlx_llama32::Llama32Runner`] with arch validation.
22
23use anyhow::{Context, Result, bail};
24use rlx_llama_base::LlamaBaseConfig;
25use std::path::{Path, PathBuf};
26
27pub use rlx_llama32::{Llama32ConfigSource, Llama32Runner, Llama32RunnerBuilder};
28
29pub const PLAN_MILESTONE: &str = "M4";
30pub const FAMILY: &str = "gpt-oss";
31
32const ACCEPTED_ARCHES: &[&str] = &["gpt-oss", "gpt_oss"];
33
34pub struct GptOssRunner {
35    inner: Llama32Runner,
36    config: LlamaBaseConfig,
37}
38
39impl GptOssRunner {
40    pub fn builder() -> GptOssRunnerBuilder {
41        GptOssRunnerBuilder::default()
42    }
43    pub fn config(&self) -> &LlamaBaseConfig {
44        &self.config
45    }
46    pub fn inner(&self) -> &Llama32Runner {
47        &self.inner
48    }
49    pub fn inner_mut(&mut self) -> &mut Llama32Runner {
50        &mut self.inner
51    }
52}
53
54#[derive(Debug, Default)]
55pub struct GptOssRunnerBuilder {
56    weights: Option<PathBuf>,
57    inner: Llama32RunnerBuilder,
58}
59
60impl GptOssRunnerBuilder {
61    pub fn weights(mut self, path: impl Into<PathBuf>) -> Self {
62        let p = path.into();
63        self.weights = Some(p.clone());
64        self.inner = self.inner.weights(p);
65        self
66    }
67    pub fn inner(mut self, f: impl FnOnce(Llama32RunnerBuilder) -> Llama32RunnerBuilder) -> Self {
68        self.inner = f(self.inner);
69        self
70    }
71
72    pub fn build(self) -> Result<GptOssRunner> {
73        let weights = self
74            .weights
75            .as_ref()
76            .ok_or_else(|| anyhow::anyhow!("weights path required (call .weights(...))"))?
77            .clone();
78        let config = LlamaBaseConfig::from_gguf_path(&weights)
79            .with_context(|| format!("rlx-gpt-oss: parse {weights:?}"))?;
80        if !ACCEPTED_ARCHES.contains(&config.arch.as_str()) {
81            bail!(
82                "rlx-gpt-oss: expected `general.architecture` ∈ {ACCEPTED_ARCHES:?}; \
83                 got `{}` at {weights:?}",
84                config.arch
85            );
86        }
87        let inner = self
88            .inner
89            .build()
90            .context("rlx-gpt-oss: building underlying Llama32Runner")?;
91        Ok(GptOssRunner { inner, config })
92    }
93}
94
95pub fn cli_run(args: &[String]) -> Result<()> {
96    if let Some(first) = args.iter().position(|a| a == "--weights") {
97        if let Some(path) = args.get(first + 1) {
98            let cfg = LlamaBaseConfig::from_gguf_path(Path::new(path))
99                .with_context(|| format!("rlx-gpt-oss: parse {path}"))?;
100            if !ACCEPTED_ARCHES.contains(&cfg.arch.as_str()) {
101                bail!(
102                    "rlx-gpt-oss: {path}: GGUF arch = `{}`, expected one of {ACCEPTED_ARCHES:?}",
103                    cfg.arch
104                );
105            }
106        }
107    }
108    rlx_llama32::cli::run(args)
109}