Skip to main content

rlx_phi/
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//! Phi 3 / Phi 4 runner.
17//!
18//! Phi-3 and Phi-4 ship as `general.architecture = phi3` in their GGUF
19//! converters (Phi-4 reuses the Phi-3 arch tag upstream — there's no
20//! separate `phi4` enum in llama.cpp). Inference is implemented in
21//! [`rlx_llama32::Llama32Runner`] with Phi-specific partial RoPE and
22//! NeoX pairing.
23
24use anyhow::{Context, Result, bail};
25use rlx_llama_base::LlamaBaseConfig;
26use rlx_runtime::Device;
27use std::path::{Path, PathBuf};
28
29pub use rlx_llama32::{
30    Llama32Config, Llama32ConfigSource, Llama32Runner, Llama32RunnerBuilder, RopeStyle,
31};
32
33pub const PLAN_MILESTONE: &str = "M4";
34pub const FAMILY: &str = "Phi 3 / Phi 4";
35
36const ACCEPTED_ARCHES: &[&str] = &["phi3", "phi4"];
37
38pub struct PhiRunner {
39    inner: Llama32Runner,
40    config: LlamaBaseConfig,
41}
42
43impl PhiRunner {
44    pub fn builder() -> PhiRunnerBuilder {
45        PhiRunnerBuilder::default()
46    }
47    pub fn config(&self) -> &LlamaBaseConfig {
48        &self.config
49    }
50    pub fn inner(&self) -> &Llama32Runner {
51        &self.inner
52    }
53    pub fn inner_mut(&mut self) -> &mut Llama32Runner {
54        &mut self.inner
55    }
56    pub fn predict_logits(&mut self, prompt_ids: &[u32]) -> Result<Vec<f32>> {
57        self.inner.predict_logits(prompt_ids)
58    }
59    pub fn generate(
60        &mut self,
61        prompt_ids: &[u32],
62        n_new: usize,
63        on_token: impl FnMut(u32),
64    ) -> Result<Vec<u32>> {
65        self.inner.generate(prompt_ids, n_new, on_token)
66    }
67    pub fn generate_packed(
68        &mut self,
69        prompt_ids: &[u32],
70        n_new: usize,
71        on_token: impl FnMut(u32),
72    ) -> Result<Vec<u32>> {
73        self.inner.generate_packed(prompt_ids, n_new, on_token)
74    }
75}
76
77#[derive(Debug, Clone, Default)]
78pub struct PhiRunnerBuilder {
79    weights: Option<PathBuf>,
80    inner: Llama32RunnerBuilder,
81}
82
83impl PhiRunnerBuilder {
84    pub fn weights(mut self, path: impl Into<PathBuf>) -> Self {
85        let p: PathBuf = path.into();
86        self.weights = Some(p.clone());
87        self.inner = self.inner.weights(p);
88        self
89    }
90    pub fn device(mut self, d: Device) -> Self {
91        self.inner = self.inner.device(d);
92        self
93    }
94    pub fn max_seq(mut self, n: usize) -> Self {
95        self.inner = self.inner.max_seq(n);
96        self
97    }
98    pub fn packed_weights(mut self, on: bool) -> Self {
99        self.inner = self.inner.packed_weights(on);
100        self
101    }
102    pub fn build(self) -> Result<PhiRunner> {
103        let weights = self
104            .weights
105            .as_ref()
106            .ok_or_else(|| anyhow::anyhow!("weights path required"))?
107            .clone();
108        let config = LlamaBaseConfig::from_gguf_path(&weights)
109            .with_context(|| format!("rlx-phi: parse {weights:?}"))?;
110        if !ACCEPTED_ARCHES.contains(&config.arch.as_str()) {
111            bail!(
112                "rlx-phi: expected `general.architecture` ∈ {ACCEPTED_ARCHES:?}; got `{}` at {weights:?}",
113                config.arch
114            );
115        }
116        let inner = self
117            .inner
118            .build()
119            .context("rlx-phi: building underlying Llama32Runner")?;
120        Ok(PhiRunner { inner, config })
121    }
122}
123
124pub fn cli_run(args: &[String]) -> Result<()> {
125    if let Some(first) = args.iter().position(|a| a == "--weights") {
126        if let Some(path) = args.get(first + 1) {
127            let cfg = LlamaBaseConfig::from_gguf_path(Path::new(path))
128                .with_context(|| format!("rlx-phi: parse {path}"))?;
129            if !ACCEPTED_ARCHES.contains(&cfg.arch.as_str()) {
130                bail!(
131                    "rlx-phi: {path}: GGUF arch = `{}`, expected one of {ACCEPTED_ARCHES:?}",
132                    cfg.arch
133                );
134            }
135        }
136    }
137    rlx_llama32::cli::run(args)
138}