Update README.md
Browse files
README.md
CHANGED
|
@@ -97,7 +97,8 @@ from transformers import AutoModelForSequenceClassification
|
|
| 97 |
import torch
|
| 98 |
|
| 99 |
model_name = "OpenBMB/MiniCPM-Reranker-Light"
|
| 100 |
-
|
|
|
|
| 101 |
# model = AutoModelForSequenceClassification.from_pretrained(model_name, trust_remote_code=True,attn_implementation="flash_attention_2", torch_dtype=torch.float16).to("cuda")
|
| 102 |
model.eval()
|
| 103 |
|
|
|
|
| 97 |
import torch
|
| 98 |
|
| 99 |
model_name = "OpenBMB/MiniCPM-Reranker-Light"
|
| 100 |
+
model = AutoModelForSequenceClassification.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.float16).to("cuda")
|
| 101 |
+
# You can also use the following code to use flash_attention_2
|
| 102 |
# model = AutoModelForSequenceClassification.from_pretrained(model_name, trust_remote_code=True,attn_implementation="flash_attention_2", torch_dtype=torch.float16).to("cuda")
|
| 103 |
model.eval()
|
| 104 |
|