From 887dd539293e925d4f8cee13ae848b43f8ac2eb0 Mon Sep 17 00:00:00 2001
From: Eric Buehler <65165915+EricLBuehler@users.noreply.github.com>
Date: Wed, 5 Jun 2024 05:33:20 -0400
Subject: [PATCH] More verbose logging during loading (#385)

* More verbose logging when loading

* More logging
---
 mistralrs-core/src/lib.rs             |  4 ++--
 mistralrs-core/src/pipeline/macros.rs | 20 +++++++++++++++-----
 2 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/mistralrs-core/src/lib.rs b/mistralrs-core/src/lib.rs
index 8132071ad..3876cfdc0 100644
--- a/mistralrs-core/src/lib.rs
+++ b/mistralrs-core/src/lib.rs
@@ -150,7 +150,7 @@ fn set_gemm_reduced_precision_f16() {
     let a = Tensor::zeros((2, 2), DType::BF16, &Device::new_cuda(0).unwrap()).unwrap();
     candle_core::cuda::set_gemm_reduced_precision_bf16(true);
     match a.matmul(&a) {
-        Ok(_) => (),
+        Ok(_) => tracing::info!("Enabling GEMM reduced precision in BF16."),
         Err(e) => {
             if format!("{e:?}").contains("CUBLAS_STATUS_NOT_SUPPORTED") {
                 tracing::info!("GEMM reduced precision in BF16 not supported.");
@@ -163,7 +163,7 @@ fn set_gemm_reduced_precision_f16() {
     let a = Tensor::zeros((2, 2), DType::F16, &Device::new_cuda(0).unwrap()).unwrap();
     candle_core::cuda::set_gemm_reduced_precision_f16(true);
     match a.matmul(&a) {
-        Ok(_) => (),
+        Ok(_) => tracing::info!("Enabling GEMM reduced precision in F16."),
         Err(e) => {
             if format!("{e:?}").contains("CUBLAS_STATUS_NOT_SUPPORTED") {
                 tracing::info!("GEMM reduced precision in F16 not supported.");
diff --git a/mistralrs-core/src/pipeline/macros.rs b/mistralrs-core/src/pipeline/macros.rs
index d3795c413..37a9ce4eb 100644
--- a/mistralrs-core/src/pipeline/macros.rs
+++ b/mistralrs-core/src/pipeline/macros.rs
@@ -92,9 +92,11 @@ macro_rules! get_paths {
             info!("Using tokenizer.json at `{p}`");
             PathBuf::from_str(p)?
         } else {
+            info!("Loading `tokenizer.json` at `{}`", $this.model_id);
             $crate::api_get_file!(api, "tokenizer.json", model_id)
         };
 
+        info!("Loading `config.json` at `{}`", $this.model_id);
         let config_filename = $crate::api_get_file!(api, "config.json", model_id);
 
         let filenames = get_model_paths(
@@ -125,6 +127,7 @@ macro_rules! get_paths {
             .collect::<Vec<_>>()
             .contains(&"generation_config.json".to_string())
         {
+            info!("Loading `generation_config.json` at `{}`", $this.model_id);
             Some($crate::api_get_file!(
                 api,
                 "generation_config.json",
@@ -138,6 +141,7 @@ macro_rules! get_paths {
             .collect::<Vec<_>>()
             .contains(&"preprocessor_config.json".to_string())
         {
+            info!("Loading `preprocessor_config.json` at `{}`", $this.model_id);
             Some($crate::api_get_file!(
                 api,
                 "preprocessor_config.json",
@@ -151,6 +155,7 @@ macro_rules! get_paths {
             .collect::<Vec<_>>()
             .contains(&"processor_config.json".to_string())
         {
+            info!("Loading `processor_config.json` at `{}`", $this.model_id);
             Some($crate::api_get_file!(
                 api,
                 "processor_config.json",
@@ -160,6 +165,7 @@ macro_rules! get_paths {
             None
         };
 
+        info!("Loading `tokenizer_config.json` at `{}`", $this.model_id);
         let template_filename = $crate::api_get_file!(api, "tokenizer_config.json", model_id);
 
         Ok(Box::new($path_name {
@@ -188,14 +194,13 @@ macro_rules! get_paths_gguf {
             .with_token(get_token($token_source)?)
             .build()?;
         let revision = $revision.unwrap_or("main".to_string());
-        let model_id_this = $this.model_id.clone().unwrap_or($this.quantized_model_id.clone());
-        let model_id_copy = model_id_this.clone();
+        let this_model_id = $this.model_id.clone().unwrap_or($this.quantized_model_id.clone());
         let api = api.repo(Repo::with_revision(
-            model_id_this.clone(),
+            this_model_id.clone(),
             RepoType::Model,
             revision.clone(),
         ));
-        let model_id = std::path::Path::new(&model_id_copy);
+        let model_id = std::path::Path::new(&this_model_id);
 
         let chat_template = if let Some(ref p) = $this.chat_template {
             if p.ends_with(".json") {
@@ -205,6 +210,7 @@ macro_rules! get_paths_gguf {
                 PathBuf::from_str("")?
             }
         } else {
+            info!("Loading `tokenizer_config.json` at `{}` because no chat template file was specified.", this_model_id);
             $crate::api_get_file!(
                 api,
                 "tokenizer_config.json",
@@ -229,7 +235,7 @@ macro_rules! get_paths_gguf {
             xlora_config,
             lora_preload_adapter_info,
         } = get_xlora_paths(
-            model_id_this,
+            this_model_id.clone(),
             &$this.xlora_model_id,
             &$token_source,
             revision.clone(),
@@ -240,6 +246,7 @@ macro_rules! get_paths_gguf {
             .collect::<Vec<_>>()
             .contains(&"generation_config.json".to_string())
         {
+            info!("Loading `generation_config.json` at `{}`", this_model_id);
             Some($crate::api_get_file!(
                 api,
                 "generation_config.json",
@@ -253,6 +260,7 @@ macro_rules! get_paths_gguf {
             .collect::<Vec<_>>()
             .contains(&"preprocessor_config.json".to_string())
         {
+            info!("Loading `preprocessor_config.json` at `{}`", this_model_id);
             Some($crate::api_get_file!(
                 api,
                 "preprocessor_config.json",
@@ -266,6 +274,7 @@ macro_rules! get_paths_gguf {
             .collect::<Vec<_>>()
             .contains(&"processor_config.json".to_string())
         {
+            info!("Loading `processor_config.json` at `{}`", this_model_id);
             Some($crate::api_get_file!(
                 api,
                 "processor_config.json",
@@ -276,6 +285,7 @@ macro_rules! get_paths_gguf {
         };
 
         let tokenizer_filename = if $this.model_id.is_some() {
+            info!("Loading `tokenizer.json` at `{}`", this_model_id);
             $crate::api_get_file!(api, "tokenizer.json", model_id)
         } else {
             PathBuf::from_str("")?