update models

Cargo.toml

@@ -20,6 +20,8 @@ pdf-extract = "0.7.7"
 candle-nn = { git = "https://github.com/huggingface/candle.git", version = "0.7.2" }
 candle-transformers = { git = "https://github.com/huggingface/candle.git", version = "0.7.2" }
 candle-core = { git = "https://github.com/huggingface/candle.git", version = "0.7.2" }
+candle-flash-attn = { git = "https://github.com/huggingface/candle.git", version = "0.7.2" }
+
 ort = {version = "=2.0.0-rc.8", features = ["cuda"]}
 strum = "0.26.1"
 strum_macros = "0.26"

rust/Cargo.toml

@@ -83,6 +83,8 @@ accelerate-src = { version = "0.3.2", optional = true }
 indicatif = "0.17.8"
 statistical = "1.0.0"
 half = "2.4.1"
+candle-flash-attn = { workspace = true, optional = true }
+
 
 [dev-dependencies]
 tempdir = "0.3.7"
@@ -97,3 +99,4 @@ accelerate = ["dep:accelerate-src", "candle-core/accelerate", "candle-nn/acceler
 cuda = ["candle-core/cuda", "candle-nn/cuda", "candle-core/cuda"]
 cudnn = ["candle-core/cudnn"]
 load-dynamic = ["ort/load-dynamic"]
+flash-attn = ["cuda", "candle-transformers/flash-attn", "dep:candle-flash-attn"]

rust/src/embeddings/local/colpali.rs

@@ -7,7 +7,7 @@ use anyhow::Error as E;
 use base64::Engine;
 use candle_core::{DType, Device, Tensor};
 use candle_nn::VarBuilder;
-use candle_transformers::models::{colpali::Model, paligemma};
+use crate::models::{colpali::Model, paligemma};
 use image::{DynamicImage, ImageFormat};
 
 use pdf2image::{Pages, RenderOptionsBuilder, PDF};

rust/src/embeddings/local/colpali_ort.rs

@@ -2,7 +2,7 @@ use std::{collections::HashMap, path::PathBuf};
 
 use anyhow::Error as E;
 use base64::Engine;
-use candle_transformers::models::paligemma;
+use crate::models::paligemma;
 use half::f16;
 use image::{DynamicImage, ImageFormat};
 use ndarray::prelude::*;

rust/src/models/colpali.rs

@@ -0,0 +1,42 @@
+use candle_core::{Module, Result, Tensor};
+use candle_nn::VarBuilder;
+
+use super::paligemma;
+use candle_nn::{linear, Linear};
+
+pub struct Model {
+    pub model: paligemma::Model,
+    pub custom_text_projection: Linear,
+}
+
+impl Model {
+    pub fn new(config: &paligemma::Config, vb: VarBuilder) -> Result<Self> {
+        let model = paligemma::Model::new(config, vb.pp("model"))?;
+        let custom_text_projection = linear(
+            config.text_config.hidden_size,
+            128,
+            vb.pp("custom_text_proj"),
+        )?;
+
+        Ok(Self {
+            model,
+            custom_text_projection,
+        })
+    }
+
+    pub fn forward_images(&mut self, pixel_values: &Tensor, input_ids: &Tensor) -> Result<Tensor> {
+        let outputs = self
+            .model
+            .setup_without_projection(pixel_values, input_ids)?;
+        let outputs = self.custom_text_projection.forward(&outputs)?;
+        let outputs = outputs.broadcast_div(&outputs.sqr()?.sum_keepdim(2)?.sqrt()?)?;
+        Ok(outputs)
+    }
+
+    pub fn forward_text(&mut self, input_ids: &Tensor) -> Result<Tensor> {
+        let outputs = self.model.forward_without_projection(input_ids)?;
+        let outputs = self.custom_text_projection.forward(&outputs)?;
+        let outputs = outputs.broadcast_div(&outputs.sqr()?.sum_keepdim(2)?.sqrt()?)?;
+        Ok(outputs)
+    }
+}