add indexer example used for benchmark

crates/core/Cargo.toml

@@ -145,3 +145,6 @@ name = "search_preindexed"
 
 [[example]]
 name = "bitvec_similarity"
+
+[[example]]
+name = "indexer"

crates/core/examples/indexer.rs

@@ -0,0 +1,62 @@
+use std::path::Path;
+
+use clap::Parser;
+use stract::config::IndexingLocalConfig;
+
+#[derive(Parser)]
+struct Args {
+    dual_encoder_path: Option<String>,
+}
+
+fn main() -> anyhow::Result<()> {
+    let args = Args::parse();
+
+    let data_path = Path::new("data");
+    let path = data_path.join("example_index");
+
+    if path.exists() {
+        std::fs::remove_dir_all(&path)?;
+    }
+
+    println!("Indexing...");
+    let start = std::time::Instant::now();
+    stract::entrypoint::indexer::run(&IndexingLocalConfig {
+        output_path: path.to_str().unwrap().to_string(),
+        limit_warc_files: None,
+        skip_warc_files: None,
+        warc_source: stract::config::WarcSource::Local(stract::config::LocalConfig {
+            folder: ".".to_string(),
+            names: vec![data_path
+                .join("sample.warc.gz")
+                .to_str()
+                .unwrap()
+                .to_string()],
+        }),
+        page_webgraph_path: Some(
+            data_path
+                .join("webgraph_page")
+                .to_str()
+                .unwrap()
+                .to_string(),
+        ),
+        host_centrality_threshold: None,
+        topics_path: None,
+        host_centrality_store_path: data_path.join("centrality/").to_str().unwrap().to_string(),
+        page_centrality_store_path: Some(
+            data_path
+                .join("centrality_page")
+                .to_str()
+                .unwrap()
+                .to_string(),
+        ),
+        safety_classifier_path: None,
+        minimum_clean_words: None,
+        batch_size: 512,
+        dual_encoder_model_path: args.dual_encoder_path,
+    })?;
+
+    println!("Indexing took {:?}", start.elapsed());
+
+    std::fs::remove_dir(path)?;
+    Ok(())
+}

crates/core/src/models/bert.rs

@@ -78,7 +78,7 @@ impl Default for Config {
             num_hidden_layers: 12,
             num_attention_heads: 12,
             intermediate_size: 3072,
-            hidden_act: HiddenAct::Gelu,
+            hidden_act: HiddenAct::GeluApproximate,
             hidden_dropout_prob: 0.1,
             max_position_embeddings: 512,
             type_vocab_size: 2,
@@ -102,7 +102,7 @@ impl Config {
             num_hidden_layers: 6,
             num_attention_heads: 12,
             intermediate_size: 1536,
-            hidden_act: HiddenAct::Gelu,
+            hidden_act: HiddenAct::GeluApproximate,
             hidden_dropout_prob: 0.1,
             max_position_embeddings: 512,
             type_vocab_size: 2,

crates/core/src/warc.rs

@@ -47,7 +47,8 @@ fn decode_string(raw: &[u8]) -> String {
         res
     } else {
         let mut detector = chardetng::EncodingDetector::new();
-        detector.feed(raw, true);
+        let end = std::cmp::min(64, raw.len());
+        detector.feed(&raw[..end], false);
         let (enc, conf) = detector.guess_assess(None, true);
 
         if conf {