Update scripts/transformers_demo.py
Browse files
scripts/transformers_demo.py
CHANGED
@@ -1,7 +1,7 @@
|
|
1 |
from transformers import AutoModelForSequenceClassification
|
2 |
import torch
|
3 |
|
4 |
-
model_name = "OpenBMB/
|
5 |
|
6 |
model = AutoModelForSequenceClassification.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.float16).to("cuda")
|
7 |
# You can use flash-attention 2 to speed up the inference
|
|
|
1 |
from transformers import AutoModelForSequenceClassification
|
2 |
import torch
|
3 |
|
4 |
+
model_name = "OpenBMB/MiniCPM-Reranker-Light"
|
5 |
|
6 |
model = AutoModelForSequenceClassification.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.float16).to("cuda")
|
7 |
# You can use flash-attention 2 to speed up the inference
|