add small modification to tests

Cargo.toml

@@ -16,10 +16,10 @@ edition = "2021"
 [workspace.dependencies]
 async-trait = "0.1.78"
 axum = "0.7.5"
-candle = { git = "https://github.com/huggingface/candle", package = "candle-core", version = "0.5.0" }
-candle-flash-attn = { git = "https://github.com/huggingface/candle", package = "candle-flash-attn", version = "0.5.0" }
-candle-nn = { git = "https://github.com/huggingface/candle", package = "candle-nn", version = "0.5.0" }
-candle-transformers = { git = "https://github.com/huggingface/candle", package = "candle-transformers", version = "0.5.0" }
+candle = { git = "https://github.com/huggingface/candle", package = "candle-core", branch = "main" }
+candle-flash-attn = { git = "https://github.com/huggingface/candle", package = "candle-flash-attn", branch = "main" }
+candle-nn = { git = "https://github.com/huggingface/candle", package = "candle-nn", branch = "main" }
+candle-transformers = { git = "https://github.com/huggingface/candle", package = "candle-transformers", branch = "main" }
 config = "0.14.0"
 dotenv = "0.15.0"
 ed25519-consensus = "2.1.0"

atoma-inference/src/models/candle/falcon.rs

@@ -278,8 +278,8 @@ mod tests {
         );
         let output = model.run(input).expect("Failed to run inference");
 
-        assert!(output.len() >= 1);
-        assert!(output.split(" ").collect::<Vec<_>>().len() <= max_tokens);
+        assert!(output.text.len() >= 1);
+        assert!(output.text.split(" ").collect::<Vec<_>>().len() <= max_tokens);
 
         std::fs::remove_dir_all(cache_dir).unwrap();
     }
@@ -361,8 +361,8 @@ mod tests {
         let output = model.run(input).expect("Failed to run inference");
         println!("{output}");
 
-        assert!(output.len() >= 1);
-        assert!(output.split(" ").collect::<Vec<_>>().len() <= max_tokens);
+        assert!(output.text.len() >= 1);
+        assert!(output.text.split(" ").collect::<Vec<_>>().len() <= max_tokens);
 
         std::fs::remove_dir_all(cache_dir).unwrap();
     }

atoma-inference/src/models/candle/llama.rs

@@ -253,7 +253,6 @@ mod tests {
             panic!("Invalid device")
         }
 
-        assert!(model.cache.use_kv_cache);
         assert_eq!(model.model_type, ModelType::LlamaTinyLlama1_1BChat);
 
         let prompt = "Write a hello world rust program: ".to_string();