temporarily add reranker
This commit is contained in:
parent
08530edb3d
commit
d1f94c2846
|
|
@ -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")
|
||||
|
||||
Loading…
Reference in New Issue