Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #25 from meilisearch/search-movies
Split the import from the searches in the movies example
- Loading branch information
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,108 +1,86 @@ | ||
| //! Download the associated file at https://www.notion.so/meilisearch/Movies-embeddings-1de3258859f54b799b7883882219d266 | ||
| use std::fs; | ||
| use std::io::{BufRead, BufReader}; | ||
| use std::path::PathBuf; | ||
| use std::time::Instant; | ||
|
|
||
| use arroy::distances::DotProduct; | ||
| use arroy::{Database, Writer}; | ||
| use clap::Parser; | ||
| use heed::EnvOpenOptions; | ||
| use rand::rngs::StdRng; | ||
| use rand::SeedableRng; | ||
|
|
||
| /// 2 GiB | ||
| const DEFAULT_MAP_SIZE: usize = 1024 * 1024 * 1024 * 2; | ||
|
|
||
| #[derive(Parser)] | ||
| #[command(author, version, about, long_about = None)] | ||
| struct Cli { | ||
| /// Sets a custom database path. | ||
| #[arg(default_value = "import.ary")] | ||
| database: PathBuf, | ||
|
|
||
| /// Specify the size of the database. | ||
| #[arg(long, default_value_t = DEFAULT_MAP_SIZE)] | ||
| map_size: usize, | ||
|
|
||
| /// The number of dimensions to construct the arroy tree. | ||
| #[arg(long)] | ||
| dimensions: usize, | ||
|
|
||
| /// The seed to generate the internal trees. | ||
| #[arg(long, default_value_t = 42)] | ||
| seed: u64, | ||
| } | ||
|
|
||
| fn main() -> Result<(), heed::BoxedError> { | ||
| let Cli { database, map_size, dimensions, seed } = Cli::parse(); | ||
|
|
||
| let mut rng = StdRng::seed_from_u64(seed); | ||
| let reader = BufReader::new(std::io::stdin()); | ||
|
|
||
| let _ = fs::create_dir_all(&database); | ||
| let env = EnvOpenOptions::new().map_size(map_size).open(&database).unwrap(); | ||
| let mut wtxn = env.write_txn().unwrap(); | ||
| let database: Database<DotProduct> = env.create_database(&mut wtxn, None)?; | ||
| let writer = Writer::<DotProduct>::prepare(&mut wtxn, database, 0, dimensions)?; | ||
|
|
||
| // The file look like that | ||
| // === BEGIN vectors === | ||
| // 0, [0.010056925, -0.0045358953, 0.009904552, 0.0046241777, ..., -0.050245073] | ||
| // === END vectors === | ||
|
|
||
| let now = Instant::now(); | ||
| let mut count = 0; | ||
| for line in reader.lines() { | ||
| let line = line?; | ||
| if line.starts_with("===") { | ||
| continue; | ||
| } | ||
|
|
||
| let (id, vector) = line.split_once(',').expect(&line); | ||
| let id: u32 = id.parse()?; | ||
| let vector: Vec<_> = vector | ||
| .trim_matches(|c: char| c.is_whitespace() || c == '[' || c == ']') | ||
| .split(',') | ||
| .map(|s| s.trim().parse::<f32>().unwrap()) | ||
| .collect(); | ||
|
|
||
| assert_eq!(vector.len(), dimensions); | ||
| writer.add_item(&mut wtxn, id, &vector)?; | ||
| count += 1; | ||
| } | ||
| println!("Took {:.2?} to parse and insert into arroy", now.elapsed()); | ||
| println!("There are {count} vectors"); | ||
| println!(); | ||
|
|
||
| println!("Building the arroy internal trees..."); | ||
| let now = Instant::now(); | ||
| writer.build(&mut wtxn, &mut rng, None)?; | ||
| wtxn.commit().unwrap(); | ||
| println!("Took {:.2?} to build", now.elapsed()); | ||
|
|
||
| Ok(()) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,53 @@ | ||
| use std::io::{BufRead, BufReader}; | ||
|
|
||
| use rand::rngs::StdRng; | ||
| use rand::seq::SliceRandom; | ||
| use rand::SeedableRng; | ||
|
|
||
| const DEFAULT_COUNT: usize = 1_000_000; | ||
|
|
||
| fn main() { | ||
| let count = std::env::args().nth(1).map(|c| c.parse().unwrap()).unwrap_or(DEFAULT_COUNT); | ||
| let reader = BufReader::new(std::io::stdin()); | ||
| let mut vectors: Vec<(u32, Vec<f32>)> = Vec::new(); | ||
|
|
||
| for line in reader.lines() { | ||
| let line = line.unwrap(); | ||
| if line.starts_with("===") { | ||
| continue; | ||
| } | ||
|
|
||
| let (id, vector) = line.split_once(',').expect(&line); | ||
| let id: u32 = id.parse().unwrap(); | ||
|
|
||
| let vector = vector | ||
| .trim_matches(|c: char| c.is_whitespace() || c == '[' || c == ']') | ||
| .split(',') | ||
| .map(|s| s.trim().parse::<f32>().unwrap()) | ||
| .collect(); | ||
|
|
||
| vectors.push((id, vector)); | ||
| assert_eq!(vectors[0].1.len(), vectors.last().unwrap().1.len()); | ||
| } | ||
|
|
||
| println!("=== BEGIN vectors ==="); | ||
| let vector_len = vectors[0].1.len(); | ||
| let mut rng = StdRng::seed_from_u64(42); | ||
| let mut id = 0; | ||
| for _ in 0..(count.checked_div(2).unwrap()) { | ||
| let mut iter = vectors.choose_multiple(&mut rng, 2); | ||
| let a = iter.next().unwrap(); | ||
| let b = iter.next().unwrap(); | ||
|
|
||
| let (al, ar) = a.1.split_at(vector_len / 2); | ||
| let (bl, br) = b.1.split_at(vector_len - al.len()); // even numbers | ||
|
|
||
| let newa: Vec<f32> = al.iter().copied().chain(br.iter().copied()).collect(); | ||
| let newb: Vec<f32> = bl.iter().copied().chain(ar.iter().copied()).collect(); | ||
| println!("{}, {:?}", id, newa); | ||
| println!("{}, {:?}", id + 1, newb); | ||
|
|
||
| id += 2; | ||
| } | ||
| println!("=== END vectors ==="); | ||
| } |
Oops, something went wrong.