|
@@ -0,0 +1,212 @@
|
|
|
|
|
+use std::{
|
|
|
|
|
+ io::{self, BufRead, Write},
|
|
|
|
|
+ process,
|
|
|
|
|
+ thread::sleep,
|
|
|
|
|
+ time::Duration,
|
|
|
|
|
+};
|
|
|
|
|
+
|
|
|
|
|
+use clap::Parser;
|
|
|
|
|
+
|
|
|
|
|
+#[derive(Parser)]
|
|
|
|
|
+#[command(about = "Binary search for best NGL and start llama-server")]
|
|
|
|
|
+struct Args {
|
|
|
|
|
+ #[arg(short, long, default_value_t = 32)]
|
|
|
|
|
+ ctx: u32,
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+ #[arg(short, long, default_value_t = 32)]
|
|
|
|
|
+ threads: u32,
|
|
|
|
|
+
|
|
|
|
|
+ #[arg(short, long, default_value_t = 11434)]
|
|
|
|
|
+ port: u16,
|
|
|
|
|
+
|
|
|
|
|
+ #[arg(long, default_value = "0.0.0.0")]
|
|
|
|
|
+ host: String,
|
|
|
|
|
+
|
|
|
|
|
+ #[arg(short, long, default_value_t = 512)]
|
|
|
|
|
+ batch: u32,
|
|
|
|
|
+
|
|
|
|
|
+ #[arg(short, long, default_value_t = 256)]
|
|
|
|
|
+ ubatch: u32,
|
|
|
|
|
+
|
|
|
|
|
+ #[arg(short)]
|
|
|
|
|
+ verbose: bool,
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+fn main() {
|
|
|
|
|
+ let args = Args::parse();
|
|
|
|
|
+ let ctx_size = args.ctx * 1024;
|
|
|
|
|
+
|
|
|
|
|
+ let model_path = select_model();
|
|
|
|
|
+
|
|
|
|
|
+ let best_fit = (u8::MIN..u8::MAX)
|
|
|
|
|
+ .collect::<Vec<_>>()
|
|
|
|
|
+ .binary_search_by(|x| {
|
|
|
|
|
+ if test_ngl(*x, &args, ctx_size, &model_path) {
|
|
|
|
|
+ std::cmp::Ordering::Less
|
|
|
|
|
+ } else {
|
|
|
|
|
+ std::cmp::Ordering::Greater
|
|
|
|
|
+ }
|
|
|
|
|
+ })
|
|
|
|
|
+ .unwrap_or_else(|x| x - 1);
|
|
|
|
|
+
|
|
|
|
|
+ eprintln!(
|
|
|
|
|
+ "final spawn with ngl={} ctx={}x1024 ({} bytes)",
|
|
|
|
|
+ best_fit,
|
|
|
|
|
+ args.ctx,
|
|
|
|
|
+ ctx_size
|
|
|
|
|
+ );
|
|
|
|
|
+
|
|
|
|
|
+ sleep(Duration::from_secs(1));
|
|
|
|
|
+
|
|
|
|
|
+ let mut child = spawn_final_server(best_fit as u8 - 1, &args, ctx_size, &model_path);
|
|
|
|
|
+
|
|
|
|
|
+ child.wait().expect("Sempre vou matar com control+c");
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+fn select_model() -> String {
|
|
|
|
|
+ let home_dir = home::home_dir().expect("home dir not found");
|
|
|
|
|
+ let model_dir = home_dir.join("guff_models");
|
|
|
|
|
+
|
|
|
|
|
+ let entries = match std::fs::read_dir(&model_dir) {
|
|
|
|
|
+ Ok(entries) => entries,
|
|
|
|
|
+ Err(e) => {
|
|
|
|
|
+ eprintln!("Error: cannot access '{}': {}", model_dir.display(), e);
|
|
|
|
|
+ std::process::exit(1);
|
|
|
|
|
+ }
|
|
|
|
|
+ };
|
|
|
|
|
+
|
|
|
|
|
+ let mut models: Vec<std::path::PathBuf> = entries
|
|
|
|
|
+ .filter_map(|e| e.ok())
|
|
|
|
|
+ .filter(|e| {
|
|
|
|
|
+ e.file_name()
|
|
|
|
|
+ .to_string_lossy()
|
|
|
|
|
+ .to_lowercase()
|
|
|
|
|
+ .ends_with(".gguf")
|
|
|
|
|
+ })
|
|
|
|
|
+ .map(|e| e.path())
|
|
|
|
|
+ .collect();
|
|
|
|
|
+
|
|
|
|
|
+ models.sort_by(|a, b| {
|
|
|
|
|
+ let a_name = a.file_name().unwrap_or_default();
|
|
|
|
|
+ let b_name = b.file_name().unwrap_or_default();
|
|
|
|
|
+ a_name.cmp(b_name)
|
|
|
|
|
+ });
|
|
|
|
|
+
|
|
|
|
|
+ if models.is_empty() {
|
|
|
|
|
+ eprintln!("No .gguf files found in {}", model_dir.display());
|
|
|
|
|
+ std::process::exit(1);
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ println!("\nAvailable models:");
|
|
|
|
|
+ for (i, path) in models.iter().enumerate() {
|
|
|
|
|
+ let label = (b'a' + i as u8) as char;
|
|
|
|
|
+ let name = path
|
|
|
|
|
+ .file_name()
|
|
|
|
|
+ .unwrap_or_default()
|
|
|
|
|
+ .to_string_lossy();
|
|
|
|
|
+ println!(" {label}) {name}");
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ loop {
|
|
|
|
|
+ let upper = (b'a' + models.len() as u8 - 1) as char;
|
|
|
|
|
+ print!("\nSelect model [a-{upper}]: ");
|
|
|
|
|
+ io::stdout().flush().unwrap();
|
|
|
|
|
+
|
|
|
|
|
+ let mut input = String::new();
|
|
|
|
|
+ io::stdin().read_line(&mut input).unwrap();
|
|
|
|
|
+
|
|
|
|
|
+ let c = input.trim().chars().next().unwrap_or('a');
|
|
|
|
|
+ let idx = match c.to_ascii_lowercase() {
|
|
|
|
|
+ c if c as u8 >= b'a' && (c as u8) < { b'a' + models.len() as u8 } => {
|
|
|
|
|
+ c as u8 - b'a'
|
|
|
|
|
+ }
|
|
|
|
|
+ _ => continue,
|
|
|
|
|
+ };
|
|
|
|
|
+
|
|
|
|
|
+ return models[idx as usize]
|
|
|
|
|
+ .to_string_lossy()
|
|
|
|
|
+ .into_owned();
|
|
|
|
|
+ }
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+fn test_ngl(ngl: u8, args: &Args, ctx_size: u32, model_path: &str) -> bool {
|
|
|
|
|
+ let mut child = spawn_llama_server(ngl, args, ctx_size, model_path);
|
|
|
|
|
+
|
|
|
|
|
+ let found = {
|
|
|
|
|
+ let handle = child.stderr.take().expect("stderr not available");
|
|
|
|
|
+ let reader = io::BufReader::new(handle);
|
|
|
|
|
+ reader
|
|
|
|
|
+ .lines()
|
|
|
|
|
+ .any(|line| {
|
|
|
|
|
+ if let Ok(l) = line {
|
|
|
|
|
+ if l.contains("You are a helpful assistant") {
|
|
|
|
|
+ return true;
|
|
|
|
|
+ }
|
|
|
|
|
+ }
|
|
|
|
|
+ false
|
|
|
|
|
+ })
|
|
|
|
|
+ };
|
|
|
|
|
+
|
|
|
|
|
+ let _ = child.kill();
|
|
|
|
|
+ let _ = child.wait();
|
|
|
|
|
+
|
|
|
|
|
+ sleep(Duration::from_secs(1));
|
|
|
|
|
+
|
|
|
|
|
+ if found {
|
|
|
|
|
+ eprintln!("{ngl} [ok]");
|
|
|
|
|
+ true
|
|
|
|
|
+ } else {
|
|
|
|
|
+ eprintln!("{ngl} [fail]");
|
|
|
|
|
+ false
|
|
|
|
|
+ }
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+fn spawn_llama_server(ngl: u8, args: &Args, ctx_size: u32, model_path: &str) -> process::Child {
|
|
|
|
|
+ process::Command::new("llama-server")
|
|
|
|
|
+ .arg("-m")
|
|
|
|
|
+ .arg(model_path)
|
|
|
|
|
+ .arg("-t")
|
|
|
|
|
+ .arg(args.threads.to_string())
|
|
|
|
|
+ .arg("-c")
|
|
|
|
|
+ .arg(ctx_size.to_string())
|
|
|
|
|
+ .arg("--port")
|
|
|
|
|
+ .arg(args.port.to_string())
|
|
|
|
|
+ .arg("--host")
|
|
|
|
|
+ .arg(&args.host)
|
|
|
|
|
+ .arg("-ngl")
|
|
|
|
|
+ .arg(&ngl.to_string())
|
|
|
|
|
+ .arg("-b")
|
|
|
|
|
+ .arg(args.batch.to_string())
|
|
|
|
|
+ .arg("-ub")
|
|
|
|
|
+ .arg(args.ubatch.to_string())
|
|
|
|
|
+ .arg("--flash-attn")
|
|
|
|
|
+ .arg("on")
|
|
|
|
|
+ .stdout(std::process::Stdio::piped())
|
|
|
|
|
+ .stderr(std::process::Stdio::piped())
|
|
|
|
|
+ .spawn()
|
|
|
|
|
+ .expect("O spawn precisa sempre acontecer")
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+ fn spawn_final_server(ngl: u8, args: &Args, ctx_size: u32, model_path: &str) -> process::Child {
|
|
|
|
|
+ let mut cmd = process::Command::new("llama-server");
|
|
|
|
|
+ cmd.arg("-m").arg(model_path)
|
|
|
|
|
+ .arg("-t").arg(args.threads.to_string())
|
|
|
|
|
+ .arg("-c").arg(ctx_size.to_string())
|
|
|
|
|
+ .arg("--port").arg(args.port.to_string())
|
|
|
|
|
+ .arg("--host").arg(&args.host)
|
|
|
|
|
+ .arg("-ngl").arg(&ngl.to_string())
|
|
|
|
|
+ .arg("-b").arg(args.batch.to_string())
|
|
|
|
|
+ .arg("-ub").arg(args.ubatch.to_string())
|
|
|
|
|
+ .arg("--flash-attn").arg("on");
|
|
|
|
|
+
|
|
|
|
|
+ if args.verbose {
|
|
|
|
|
+ cmd.stdout(std::process::Stdio::inherit())
|
|
|
|
|
+ .stderr(std::process::Stdio::inherit());
|
|
|
|
|
+ } else {
|
|
|
|
|
+ cmd.stdout(std::process::Stdio::null())
|
|
|
|
|
+ .stderr(std::process::Stdio::null());
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ cmd.spawn().expect("O spawn precisa sempre acontecer")
|
|
|
|
|
+}
|