|
1 | | -// mod common; |
2 | | -// |
3 | | -// use crate::common::{start_server, Score}; |
4 | | -// use anyhow::Result; |
5 | | -// use insta::internals::YamlMatcher; |
6 | | -// use serde_json::json; |
7 | | -// use text_embeddings_backend::DType; |
8 | | -// |
9 | | -// #[tokio::test] |
10 | | -// #[cfg(feature = "http")] |
11 | | -// async fn test_embeddings() -> Result<()> { |
12 | | -// start_server( |
13 | | -// "sentence-transformers/all-MiniLM-L6-v2".to_string(), |
14 | | -// None, |
15 | | -// DType::Float32, |
16 | | -// ) |
17 | | -// .await?; |
18 | | -// |
19 | | -// let request = json!({ |
20 | | -// "inputs": "test" |
21 | | -// }); |
22 | | -// let client = reqwest::Client::new(); |
23 | | -// let res = client |
24 | | -// .post("http://0.0.0.0:8090/embed") |
25 | | -// .json(&request) |
26 | | -// .send() |
27 | | -// .await?; |
28 | | -// |
29 | | -// let embeddings_single = res.json::<Vec<Vec<Score>>>().await?; |
30 | | -// let matcher = YamlMatcher::<Vec<Vec<Score>>>::new(); |
31 | | -// insta::assert_yaml_snapshot!("embeddings_single", embeddings_single, &matcher); |
32 | | -// |
33 | | -// let test_tokens = vec![[101, 3231, 102]]; // tokenized "test" |
34 | | -// let request = json!({"inputs": &test_tokens}); |
35 | | -// let res = client |
36 | | -// .post("http://0.0.0.0:8090/embed") |
37 | | -// .json(&request) |
38 | | -// .send() |
39 | | -// .await?; |
40 | | -// |
41 | | -// let embeddings_single = res.json::<Vec<Vec<Score>>>().await?; |
42 | | -// let matcher = YamlMatcher::<Vec<Vec<Score>>>::new(); |
43 | | -// insta::assert_yaml_snapshot!("embeddings_single", embeddings_single, &matcher); |
44 | | -// |
45 | | -// let request = json!({ |
46 | | -// "inputs": vec!["test", "test", "test", "test", "test"], |
47 | | -// }); |
48 | | -// |
49 | | -// let client = reqwest::Client::new(); |
50 | | -// let res = client |
51 | | -// .post("http://0.0.0.0:8090/embed") |
52 | | -// .json(&request) |
53 | | -// .send() |
54 | | -// .await?; |
55 | | -// let embeddings_batch = res.json::<Vec<Vec<Score>>>().await?; |
56 | | -// insta::assert_yaml_snapshot!("embeddings_batch", embeddings_batch, &matcher); |
57 | | -// for embeddings in &embeddings_batch { |
58 | | -// assert_eq!(embeddings, &embeddings_single[0]); |
59 | | -// } |
60 | | -// |
61 | | -// let request = |
62 | | -// json!({"inputs": &test_tokens.repeat(request["inputs"].as_array().unwrap().len())}); |
63 | | -// let res = client |
64 | | -// .post("http://0.0.0.0:8090/embed") |
65 | | -// .json(&request) |
66 | | -// .send() |
67 | | -// .await?; |
68 | | -// |
69 | | -// let embeddings_batch = res.json::<Vec<Vec<Score>>>().await?; |
70 | | -// insta::assert_yaml_snapshot!("embeddings_batch", embeddings_batch, &matcher); |
71 | | -// for embeddings in &embeddings_batch { |
72 | | -// assert_eq!(embeddings, &embeddings_single[0]); |
73 | | -// } |
74 | | -// |
75 | | -// let request = json!({ |
76 | | -// "inputs": "test" |
77 | | -// }); |
78 | | -// |
79 | | -// let client = reqwest::Client::new(); |
80 | | -// let res = client |
81 | | -// .post("http://0.0.0.0:8090/embed_all") |
82 | | -// .json(&request) |
83 | | -// .send() |
84 | | -// .await?; |
85 | | -// |
86 | | -// let embeddings_raw = res.json::<Vec<Vec<Vec<Score>>>>().await?; |
87 | | -// let matcher = YamlMatcher::<Vec<Vec<Vec<Score>>>>::new(); |
88 | | -// insta::assert_yaml_snapshot!("embeddings_raw", embeddings_raw, &matcher); |
89 | | -// |
90 | | -// Ok(()) |
91 | | -// } |
| 1 | +mod common; |
| 2 | + |
| 3 | +use crate::common::{start_server, Score}; |
| 4 | +use anyhow::Result; |
| 5 | +use insta::internals::YamlMatcher; |
| 6 | +use serde_json::json; |
| 7 | +use text_embeddings_backend::DType; |
| 8 | + |
| 9 | +#[tokio::test] |
| 10 | +#[cfg(feature = "http")] |
| 11 | +async fn test_mrl_embeddings() -> Result<()> { |
| 12 | + start_server( |
| 13 | + "sentence-transformers/all-MiniLM-L6-v2".to_string(), |
| 14 | + None, |
| 15 | + DType::Float32, |
| 16 | + ) |
| 17 | + .await?; |
| 18 | + |
| 19 | + let request = json!({ |
| 20 | + "inputs": "test", |
| 21 | + "dimensions": 128, |
| 22 | + }); |
| 23 | + |
| 24 | + let client = reqwest::Client::new(); |
| 25 | + let res = client |
| 26 | + .post("http://0.0.0.0:8090/embed") |
| 27 | + .json(&request) |
| 28 | + .send() |
| 29 | + .await?; |
| 30 | + |
| 31 | + let embeddings_single = res.json::<Vec<Vec<Score>>>().await?; |
| 32 | + let matcher = YamlMatcher::<Vec<Vec<Score>>>::new(); |
| 33 | + insta::assert_yaml_snapshot!("mrl_embeddings_single", embeddings_single, &matcher); |
| 34 | + |
| 35 | + let request = json!({ |
| 36 | + "inputs": vec!["test", "test", "test"], |
| 37 | + "dimensions": 128, |
| 38 | + }); |
| 39 | + |
| 40 | + let client = reqwest::Client::new(); |
| 41 | + let res = client |
| 42 | + .post("http://0.0.0.0:8090/embed") |
| 43 | + .json(&request) |
| 44 | + .send() |
| 45 | + .await?; |
| 46 | + |
| 47 | + let embeddings_batch = res.json::<Vec<Vec<Score>>>().await?; |
| 48 | + insta::assert_yaml_snapshot!("mrl_embeddings_batch", embeddings_batch, &matcher); |
| 49 | + |
| 50 | + for embeddings in &embeddings_batch { |
| 51 | + assert_eq!(embeddings, &embeddings_single[0]); |
| 52 | + } |
| 53 | + |
| 54 | + Ok(()) |
| 55 | +} |
0 commit comments