StableLM -> NeoX

.vscode/launch.json

@@ -9,11 +9,7 @@
       "request": "launch",
       "name": "Debug BLOOM Inference",
       "cargo": {
-        "args": [
-          "build",
-          "--example=bloom-inference",
-          "--package=llm-bloom"
-        ],
+        "args": ["build", "--example=bloom-inference", "--package=llm-bloom"],
         "filter": {
           "name": "bloom-inference",
           "kind": "example"
@@ -27,11 +23,7 @@
       "request": "launch",
       "name": "Debug GPT-2 Inference",
       "cargo": {
-        "args": [
-          "build",
-          "--example=gpt2-inference",
-          "--package=llm-gpt2"
-        ],
+        "args": ["build", "--example=gpt2-inference", "--package=llm-gpt2"],
         "filter": {
           "name": "gpt2-inference",
           "kind": "example"
@@ -45,11 +37,7 @@
       "request": "launch",
       "name": "Debug LLaMA Inference",
       "cargo": {
-        "args": [
-          "build",
-          "--example=llama-inference",
-          "--package=llm-llama"
-        ],
+        "args": ["build", "--example=llama-inference", "--package=llm-llama"],
         "filter": {
           "name": "llama-inference",
           "kind": "example"
@@ -61,20 +49,16 @@
     {
       "type": "lldb",
       "request": "launch",
-      "name": "Debug StableLM Inference",
+      "name": "Debug NeoX Inference",
       "cargo": {
-        "args": [
-          "build",
-          "--example=stablelm-inference",
-          "--package=llm-stablelm"
-        ],
+        "args": ["build", "--example=neox-inference", "--package=llm-neox"],
         "filter": {
-          "name": "stablelm-inference",
+          "name": "neox-inference",
           "kind": "example"
         }
       },
       "args": ["${env:HOME}/.ggml-models/stablelm-base-alpha-3b-f16.bin"],
       "cwd": "${workspaceFolder}"
     }
   ]
-}
+}

crates/models/stablelm/Cargo.toml → crates/models/neox/Cargo.toml

@@ -1,5 +1,5 @@
 [package]
-name = "llm-stablelm"
+name = "llm-neox"
 version = { workspace = true }
 edition = "2021"
 

crates/models/stablelm/examples/stablelm-inference.rs → crates/models/neox/examples/neox-inference.rs

@@ -13,9 +13,8 @@ fn main() {
     println!(" >>> Loading model from {loc}...");
     let now = std::time::Instant::now();
 
-    let gpt2 =
-        llm_stablelm::StableLm::load(Path::new(loc), true, 512, load_progress_callback_stdout)
-            .unwrap_or_else(|e| panic!("Error loading model from {loc}: {e}"));
+    let gpt2 = llm_neox::NeoX::load(Path::new(loc), true, 512, load_progress_callback_stdout)
+        .unwrap_or_else(|e| panic!("Error loading model from {loc}: {e}"));
 
     println!(" >>> Model loaded in {} ms.", now.elapsed().as_millis());
 

crates/models/stablelm/src/lib.rs → crates/models/neox/src/lib.rs

@@ -9,7 +9,7 @@ use llm_base::{
     Vocabulary,
 };
 
-pub struct StableLm {
+pub struct NeoX {
     hyperparameters: Hyperparameters,
     n_context_tokens: usize,
 
@@ -34,10 +34,10 @@ pub struct StableLm {
     _context: ggml::Context,
 }
 
-unsafe impl Send for StableLm {}
-unsafe impl Sync for StableLm {}
+unsafe impl Send for NeoX {}
+unsafe impl Sync for NeoX {}
 
-impl StableLm {
+impl NeoX {
     /// Load the model from `path` with `n_context_tokens` context tokens.
     ///
     /// The status of the loading process will be reported through `load_progress_callback`.
@@ -46,12 +46,12 @@ impl StableLm {
         prefer_mmap: bool,
         n_context_tokens: usize,
         load_progress_callback: impl FnMut(LoadProgress),
-    ) -> Result<StableLm, LoadError> {
+    ) -> Result<NeoX, LoadError> {
         llm_base::load(path, prefer_mmap, n_context_tokens, load_progress_callback)
     }
 }
 
-impl KnownModel for StableLm {
+impl KnownModel for NeoX {
     type Hyperparameters = Hyperparameters;
 
     fn new<E: Error>(
@@ -138,7 +138,7 @@ impl KnownModel for StableLm {
 
         let (_context, _, _mmap) = tl.finish();
 
-        Ok(StableLm {
+        Ok(NeoX {
             hyperparameters,
             n_context_tokens,
             vocabulary,
@@ -494,7 +494,7 @@ struct Layer {
 }
 
 #[cfg(test)]
-impl StableLm {
+impl NeoX {
     /// This does *not* construct a valid model. All of the tensors are entirely
     /// empty. However, it can be used to determine if some code will compile.
     fn new_empty() -> Self {
@@ -522,7 +522,7 @@ mod tests {
 
     #[test]
     fn can_share_model_between_threads() {
-        let model = Arc::new(StableLm::new_empty());
+        let model = Arc::new(NeoX::new_empty());
 
         for _ in 0..4 {
             let model = model.clone();