Merge pull request #36 from setzer22/feat/seed_rng

Adds a way to specify the seed for RNG

Author: setzer22

llama-cli/src/cli_args.rs

@@ -67,6 +67,12 @@ pub struct Args {
     /// --cache-prompt
     #[arg(long, default_value = None)]
     pub restore_prompt: Option<String>,
+
+    /// Specifies the seed to use during sampling. Note that, depending on
+    /// hardware, the same seed may lead to different results on two separate
+    /// machines.
+    #[arg(long, default_value = None)]
+    pub seed: Option<u64>,
 }
 
 /// CLI args are stored in a lazy static variable so they're accessible from

llama-cli/src/main.rs

@@ -2,7 +2,7 @@ use std::{convert::Infallible, io::Write};
 
 use cli_args::CLI_ARGS;
 use llama_rs::{InferenceParameters, InferenceSnapshot};
-use rand::thread_rng;
+use rand::SeedableRng;
 
 mod cli_args;
 
@@ -94,7 +94,11 @@ fn main() {
 
     log::info!("Model fully loaded!");
 
-    let mut rng = thread_rng();
+    let mut rng = if let Some(seed) = CLI_ARGS.seed {
+        rand::rngs::StdRng::seed_from_u64(seed)
+    } else {
+        rand::rngs::StdRng::from_entropy()
+    };
 
     let mut session = if let Some(restore_path) = &args.restore_prompt {
         let snapshot = InferenceSnapshot::load_from_disk(restore_path);