update reranker v2

545999961 · 545999961 · commit 282776f7ebd3 · 2024-03-19T06:18:12.000Z
diff --git a/FlagEmbedding/llm_reranker/README.md b/FlagEmbedding/llm_reranker/README.md
@@ -221,7 +221,7 @@ def get_inputs(pairs, tokenizer, prompt=None, max_length=1024):
     )
 
 tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-reranker-v2-minicpm-layerwise', trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained('BAAI/bge-reranker-v2-minicpm-layerwise', trust_remote_code=True, torch_dtype=torch.float16)
+model = AutoModelForCausalLM.from_pretrained('BAAI/bge-reranker-v2-minicpm-layerwise', trust_remote_code=True, torch_dtype=torch.bfloat16)
 model = model.to('cuda')
 model.eval()
 

Original file line number	Diff line number	Diff line change
`@@ -221,7 +221,7 @@ def get_inputs(pairs, tokenizer, prompt=None, max_length=1024):`
`221`	`221`	`)`
`222`	`222`
`223`	`223`	`tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-reranker-v2-minicpm-layerwise', trust_remote_code=True)`
`224`		`-model = AutoModelForCausalLM.from_pretrained('BAAI/bge-reranker-v2-minicpm-layerwise', trust_remote_code=True, torch_dtype=torch.float16)`
	`224`	`+model = AutoModelForCausalLM.from_pretrained('BAAI/bge-reranker-v2-minicpm-layerwise', trust_remote_code=True, torch_dtype=torch.bfloat16)`
`225`	`225`	`model = model.to('cuda')`
`226`	`226`	`model.eval()`
`227`	`227`