Retrieve the topn most relevant documents for the given query. Parameters: - query (str): The input query. - docs (list of str): The list of documents to search from. - model_name (str): The name of the SentenceTransformer model to use. - width (int): The number of top docu
(query, docs, model, width=3)
| 9 | from sentence_transformers import SentenceTransformer |
| 10 | |
| 11 | def retrieve_top_docs(query, docs, model, width=3): |
| 12 | """ |
| 13 | Retrieve the topn most relevant documents for the given query. |
| 14 | |
| 15 | Parameters: |
| 16 | - query (str): The input query. |
| 17 | - docs (list of str): The list of documents to search from. |
| 18 | - model_name (str): The name of the SentenceTransformer model to use. |
| 19 | - width (int): The number of top documents to return. |
| 20 | |
| 21 | Returns: |
| 22 | - list of float: A list of scores for the topn documents. |
| 23 | - list of str: A list of the topn documents. |
| 24 | """ |
| 25 | |
| 26 | query_emb = model.encode(query) |
| 27 | doc_emb = model.encode(docs) |
| 28 | |
| 29 | scores = util.dot_score(query_emb, doc_emb)[0].cpu().tolist() |
| 30 | |
| 31 | doc_score_pairs = sorted(list(zip(docs, scores)), key=lambda x: x[1], reverse=True) |
| 32 | |
| 33 | top_docs = [pair[0] for pair in doc_score_pairs[:width]] |
| 34 | top_scores = [pair[1] for pair in doc_score_pairs[:width]] |
| 35 | |
| 36 | return top_docs, top_scores |
| 37 | |
| 38 | |
| 39 | def compute_bm25_similarity(query, corpus, width=3): |
no outgoing calls
no test coverage detected