1use 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}