Commit d5d48338 by tinywell

feat: Add CrossEncoder model to BgeRerank class

parent b538b447
......@@ -13,13 +13,12 @@ class BgeRerank(BaseDocumentCompressor):
"""Model name to use for reranking."""
top_n: int = 10
"""Number of documents to return."""
model:CrossEncoder
model:CrossEncoder = None
"""CrossEncoder instance to use for reranking."""
def __init__(self, model_name: str, top_n: int = 10):
self.model_name = model_name
super().__init__(model_name=model_name, top_n=top_n)
self.model = CrossEncoder(model_name)
self.top_n = top_n
def bge_rerank(self,query,docs):
model_inputs = [[query, doc] for doc in docs]
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment