Initial public release of Warp.

Repo-Sync-Origin: warpdotdev/warp-internal@12af1d983b

crates/input_classifier/src/onnx/candle.rs

@@ -0,0 +1,101 @@
+use std::collections::HashMap;
+
+use anyhow::{Context as _, Result, ensure};
+use candle_core::{IndexOp as _, Tensor};
+use candle_onnx::onnx::ModelProto;
+use prost::Message as _;
+use tokenizers::Tokenizer;
+use warp_completer::ParsedTokensSnapshot;
+
+use super::ClassificationResult;
+
+use super::Model;
+
+pub struct InferenceRunner {
+    model: ModelProto,
+    tokenizer: Tokenizer,
+}
+
+impl InferenceRunner {
+    pub fn new(model: Model) -> Result<Self> {
+        Ok(Self {
+            model: Self::load_model(model)?,
+            tokenizer: Self::load_tokenizer(model)?,
+        })
+    }
+
+    fn load_model(model: Model) -> Result<ModelProto> {
+        let model_bytes = model.bytes().ok_or_else(|| {
+            std::io::Error::new(std::io::ErrorKind::NotFound, "Model file not found")
+        })?;
+        let model = ModelProto::decode(model_bytes.as_ref())?;
+        Ok(model)
+    }
+
+    fn load_tokenizer(model: Model) -> Result<Tokenizer> {
+        let tokenizer_bytes = model.tokenizer_bytes().ok_or_else(|| {
+            std::io::Error::new(std::io::ErrorKind::NotFound, "Tokenizer file not found")
+        })?;
+        let tokenizer = Tokenizer::from_bytes(tokenizer_bytes).map_err(|e| anyhow::anyhow!(e))?;
+        Ok(tokenizer)
+    }
+}
+
+impl super::InferenceRunner for InferenceRunner {
+    fn run_inference(&self, input: &ParsedTokensSnapshot) -> Result<ClassificationResult> {
+        // Encode the input text into tokens.
+        let encoding = self
+            .tokenizer
+            .encode_fast(input.buffer_text.as_str(), true)
+            .map_err(|e| anyhow::anyhow!(e))?;
+
+        // For now, we'll do all inference on the CPU.
+        let device = candle_core::Device::Cpu;
+
+        let input_ids = Tensor::new(
+            encoding
+                .get_ids()
+                .iter()
+                .map(|&x| x as i64)
+                .collect::<Vec<_>>()
+                .as_slice(),
+            &device,
+        )
+        .context("failed to build input ids tensor")?;
+        let attention_mask = Tensor::new(
+            encoding
+                .get_attention_mask()
+                .iter()
+                .map(|&x| x as i64)
+                .collect::<Vec<_>>()
+                .as_slice(),
+            &device,
+        )
+        .context("failed to build attention mask tensor")?;
+
+        // Run inference.
+        let outputs = candle_onnx::simple_eval(
+            &self.model,
+            HashMap::from([
+                ("input_ids".to_string(), input_ids.unsqueeze(0)?),
+                ("attention_mask".to_string(), attention_mask.unsqueeze(0)?),
+            ]),
+        )
+        .context("error evaluating the model")?;
+
+        let logits = outputs.get("logits").context("failed to get logits")?;
+        let probabilities = candle_nn::ops::softmax_last_dim(logits)
+            .context("failed to compute softmax")?
+            .i(0)
+            .context("failed to get first dimension")?
+            .to_vec1::<f32>()
+            .context("failed to convert softmax output to vec")?;
+
+        ensure!(probabilities.len() == 2, "expected 2 probabilities");
+
+        Ok(ClassificationResult {
+            p_ai: probabilities[0],
+            p_shell: probabilities[1],
+        })
+    }
+}

crates/input_classifier/src/onnx/mod.rs

@@ -0,0 +1,198 @@
+#[cfg(feature = "onnx_candle")]
+mod candle;
+#[cfg(feature = "onnx_ort")]
+mod ort;
+
+use std::borrow::Cow;
+
+use anyhow::Result;
+use async_trait::async_trait;
+use rust_embed::RustEmbed;
+use warp_completer::ParsedTokensSnapshot;
+
+use crate::{
+    ClassificationResult, Context, InputClassifier, InputType,
+    parser::parse_query_into_tokens,
+    util::{
+        is_likely_shell_command, is_one_off_natural_language_word, is_one_off_shell_command_keyword,
+    },
+};
+
+#[derive(Clone, Copy, RustEmbed)]
+#[folder = "models/onnx"]
+struct Models;
+
+#[derive(Copy, Clone)]
+pub enum Model {
+    BertTiny,
+}
+
+impl Model {
+    fn bytes(&self) -> Option<Cow<'static, [u8]>> {
+        Models::get(self.model_path()).map(|file| file.data)
+    }
+
+    fn tokenizer_bytes(&self) -> Option<Cow<'static, [u8]>> {
+        Models::get(self.tokenizer_path()).map(|file| file.data)
+    }
+
+    fn model_path(&self) -> &'static str {
+        match self {
+            Model::BertTiny => "bert_tiny.onnx",
+        }
+    }
+
+    fn tokenizer_path(&self) -> &'static str {
+        match self {
+            Model::BertTiny => "bert_tiny_tokenizer.json",
+        }
+    }
+}
+
+pub struct OnnxClassifier {
+    inference_runner: Box<dyn InferenceRunner>,
+    has_panicked: HasPanicked,
+}
+
+impl OnnxClassifier {
+    pub fn new(_model: Model) -> Result<Self> {
+        #[cfg(feature = "onnx_candle")]
+        match candle::InferenceRunner::new(_model).map(Box::new) {
+            Ok(inference_runner) => {
+                return Ok(Self {
+                    inference_runner,
+                    has_panicked: HasPanicked::new(),
+                });
+            }
+            Err(err) => log::warn!("Failed to initialize candle inference runner: {err:#}"),
+        }
+
+        #[cfg(feature = "onnx_ort")]
+        match ort::InferenceRunner::new(_model).map(Box::new) {
+            Ok(inference_runner) => {
+                return Ok(Self {
+                    inference_runner,
+                    has_panicked: HasPanicked::new(),
+                });
+            }
+            Err(err) => log::warn!("Failed to initialize ort inference runner: {err:#}"),
+        }
+
+        Err(anyhow::anyhow!("No onnx inference engine enabled"))
+    }
+}
+
+#[cfg_attr(not(target_family = "wasm"), async_trait)]
+#[cfg_attr(target_family = "wasm", async_trait(?Send))]
+impl InputClassifier for OnnxClassifier {
+    async fn detect_input_type(&self, input: ParsedTokensSnapshot, context: &Context) -> InputType {
+        let word_tokens = parse_query_into_tokens(input.buffer_text.as_str());
+
+        let total_word_token_count = word_tokens.len();
+
+        // Start by applying some simple heuristics before running the full classifier.
+        if let Some(first_word) = word_tokens.first() {
+            let first_word = first_word.to_lowercase();
+
+            // If the input is a single word and the word is one of a specific set of words, classify it as AI
+            if word_tokens.len() == 1 && is_one_off_natural_language_word(&first_word) {
+                return InputType::AI;
+            }
+
+            // If the first token is one of a specific set of shell command keywords (e.g.: echo or sudo),
+            // we should classify it as shell.
+            if is_one_off_shell_command_keyword(&first_word) {
+                return InputType::Shell;
+            }
+        }
+
+        if is_likely_shell_command(&input, total_word_token_count).await {
+            return InputType::Shell;
+        }
+
+        // Otherwise, defer all decision-making to the model.
+        self.classify_input(input, context)
+            .await
+            .map(|result| result.to_input_type())
+            .unwrap_or(context.current_input_type)
+    }
+
+    async fn classify_input(
+        &self,
+        input: warp_completer::ParsedTokensSnapshot,
+        _context: &Context,
+    ) -> anyhow::Result<ClassificationResult> {
+        // If we ever panicked while running inference, we should fall back to the heuristic classifier.
+        if self.has_panicked.has_panicked() {
+            return crate::heuristic_classifier::HeuristicClassifier
+                .classify_input(input, _context)
+                .await;
+        }
+
+        // Given that we only can get here if we have never panicked, we don't have to
+        // worry about attempting to use an inference runner that is in an invalid state
+        // due to recovering after catching a panic unwind.
+        let inference_runner = std::panic::AssertUnwindSafe(&self.inference_runner);
+
+        let input_ref = &input;
+        match std::panic::catch_unwind(move || {
+            let start = instant::Instant::now();
+            let result = inference_runner.run_inference(input_ref);
+            let duration = start.elapsed();
+            let duration_ms = duration.as_secs_f32() * 1000.0;
+
+            match result {
+                Ok(result) => {
+                    log::debug!(
+                        "Inference took {duration_ms:.2} ms; p_shell: {:.5}, p_ai: {:.5}",
+                        result.p_shell,
+                        result.p_ai
+                    );
+                    Ok(result)
+                }
+                Err(e) => {
+                    log::error!("Failed to run inference (took {duration_ms:.2} ms): {e:#}");
+                    Err(e)
+                }
+            }
+        }) {
+            Ok(result) => result,
+            Err(_) => {
+                log::error!(
+                    "Caught panic while running inference; falling back to heuristic classifier."
+                );
+                self.has_panicked.on_panic();
+                crate::heuristic_classifier::HeuristicClassifier
+                    .classify_input(input, _context)
+                    .await
+            }
+        }
+    }
+}
+
+trait InferenceRunner: 'static + Send + Sync {
+    fn run_inference(&self, input: &ParsedTokensSnapshot) -> Result<ClassificationResult>;
+}
+
+/// A simple structure that we can use to track whether the ONNX classifier has panicked.
+struct HasPanicked {
+    inner: std::sync::Once,
+}
+
+impl HasPanicked {
+    fn new() -> Self {
+        Self {
+            inner: std::sync::Once::new(),
+        }
+    }
+
+    fn on_panic(&self) {
+        // Mark the classifier as having panicked.
+        self.inner.call_once(|| {});
+    }
+
+    fn has_panicked(&self) -> bool {
+        // Return true if the classifier has panicked.
+        self.inner.is_completed()
+    }
+}

crates/input_classifier/src/onnx/ort.rs

@@ -0,0 +1,91 @@
+use anyhow::{Result, ensure};
+use itertools::Itertools as _;
+use ort::{
+    execution_providers::CPUExecutionProvider, session::Session, tensor::ArrayExtensions as _,
+    value::Value,
+};
+use parking_lot::Mutex;
+use tokenizers::Tokenizer;
+use warp_completer::ParsedTokensSnapshot;
+
+use super::ClassificationResult;
+
+use super::Model;
+
+pub struct InferenceRunner {
+    session: Mutex<Session>,
+    tokenizer: Tokenizer,
+}
+
+impl InferenceRunner {
+    pub fn new(model: Model) -> Result<Self> {
+        Ok(Self {
+            session: Self::init_session(model)?.into(),
+            tokenizer: Self::load_tokenizer(model)?,
+        })
+    }
+
+    fn init_session(model: Model) -> Result<Session> {
+        let model_bytes = model.bytes().ok_or_else(|| {
+            std::io::Error::new(std::io::ErrorKind::NotFound, "Model file not found")
+        })?;
+        let session = Session::builder()?
+            // For now, we'll do all inference on the CPU.
+            .with_execution_providers([CPUExecutionProvider::default().build()])?
+            .commit_from_memory(model_bytes.as_ref())?;
+        Ok(session)
+    }
+
+    fn load_tokenizer(model: Model) -> Result<Tokenizer> {
+        let tokenizer_bytes = model.tokenizer_bytes().ok_or_else(|| {
+            std::io::Error::new(std::io::ErrorKind::NotFound, "Tokenizer file not found")
+        })?;
+        let tokenizer = Tokenizer::from_bytes(tokenizer_bytes).map_err(|e| anyhow::anyhow!(e))?;
+        Ok(tokenizer)
+    }
+}
+
+impl super::InferenceRunner for InferenceRunner {
+    fn run_inference(&self, input: &ParsedTokensSnapshot) -> Result<ClassificationResult> {
+        // Encode the input text into tokens.
+        let encoding = self
+            .tokenizer
+            .encode_fast(input.buffer_text.as_str(), true)
+            .map_err(|e| anyhow::anyhow!(e))?;
+
+        let input_ids = encoding.get_ids();
+        let attention_mask = encoding.get_attention_mask();
+
+        let input_ids = Value::from_array((
+            [1, input_ids.len()],
+            input_ids.iter().map(|&x| x as i64).collect_vec(),
+        ))?;
+        let attention_mask = Value::from_array((
+            [1, attention_mask.len()],
+            attention_mask.iter().map(|&x| x as i64).collect_vec(),
+        ))?;
+
+        let mut session = self.session.lock();
+        let outputs = session.run(ort::inputs![
+            "input_ids" => input_ids,
+            "attention_mask" => attention_mask,
+        ])?;
+
+        let logits = &outputs[0];
+
+        let logits = logits.try_extract_array::<f32>()?;
+        let probabilities = logits.softmax(ndarray::Axis(1));
+
+        let probabilities = probabilities.view();
+        let probabilities = probabilities
+            .as_slice()
+            .ok_or_else(|| anyhow::anyhow!("failed to get probabilities"))?;
+
+        ensure!(probabilities.len() == 2, "expected 2 probabilities");
+
+        Ok(ClassificationResult {
+            p_ai: probabilities[0],
+            p_shell: probabilities[1],
+        })
+    }
+}