Skip to content

Instantly share code, notes, and snippets.

@a-agmon
Created December 21, 2023 09:23
Show Gist options
  • Save a-agmon/45a28452ef77bc875f35783a634da1b4 to your computer and use it in GitHub Desktop.
Save a-agmon/45a28452ef77bc875f35783a634da1b4 to your computer and use it in GitHub Desktop.
find similar function
async fn find_similar(
State(model_ctx): State<Arc<(BertInferenceModel, Vec<String>)>>,
Json(payload): Json<ReqPayload>,
) -> Json<ResPayload> {
let (model, text_map) = &*model_ctx;
let query_vector = model
.infer_sentence_embedding(&payload.text)
.expect("error infering sentence embedding");
let results: Vec<(usize, f32)> = model
.score_vector_similarity(
query_vector,
payload.num_results as usize,
)
.unwrap();
let results: Vec<String> = results
.into_iter()
.map(|r| {
let top_item_text = text_map.get(r.0).unwrap();
format!(
"Item:{} (index: {} score:{:?})",
top_item_text, r.0, r.1
)
})
.collect();
Json(ResPayload { text: results })
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment