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