Skip to content

Commit

Permalink
temporarily add reranker
Browse files Browse the repository at this point in the history
  • Loading branch information
hzg0601 committed Dec 6, 2023
1 parent 08530ed commit d1f94c2
Showing 1 changed file with 19 additions and 0 deletions.
19 changes: 19 additions & 0 deletions server/reranker/reranker.py
@@ -0,0 +1,19 @@
from langchain.retrievers.document_compressors import CohereRerank
from llama_index.postprocessor import SentenceTransformerRerank
from sentence_transformers import SentenceTransformer,CrossEncoder

model_path = "/root/autodl-tmp/models/bge-reranker-large/"
instruction = "为这个句子生成表示以用于检索相关文章:"
reranker = SentenceTransformerRerank(
top_n=5,
model="local:"+model_path,
)

reranker_model = SentenceTransformer(model_name_or_path=model_path,device="cuda")

reranker_ce = CrossEncoder(model_name=model_path,device="cuda",max_length=1024)

reranker_ce.predict([[],[]])

print("Load reranker")

0 comments on commit d1f94c2

Please sign in to comment.