Spaces:
Running
Running
ThanaritKanjanametawat
commited on
Commit
•
825e56e
1
Parent(s):
81ec3b4
Move everything to CPU
Browse files- ModelDriver.py +2 -2
ModelDriver.py
CHANGED
|
@@ -61,7 +61,7 @@ def RobertaSentinelCSAbstractInference(input_text):
|
|
| 61 |
def RobertaClassifierOpenGPTInference(input_text):
|
| 62 |
tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
|
| 63 |
model_path = "ClassifierCheckpoint/RobertaClassifierOpenGPT.pth"
|
| 64 |
-
model = RobertaForSequenceClassification.from_pretrained('roberta-base', num_labels=2)
|
| 65 |
model.load_state_dict(torch.load(model_path))
|
| 66 |
model = model.to(torch.device('cpu'))
|
| 67 |
model.eval()
|
|
@@ -83,7 +83,7 @@ def RobertaClassifierOpenGPTInference(input_text):
|
|
| 83 |
def RobertaClassifierCSAbstractInference(input_text):
|
| 84 |
tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
|
| 85 |
model_path = "ClassifierCheckpoint/RobertaClassifierCSAbstract.pth"
|
| 86 |
-
model = RobertaForSequenceClassification.from_pretrained('roberta-base', num_labels=2)
|
| 87 |
model.load_state_dict(torch.load(model_path))
|
| 88 |
model = model.to(torch.device('cpu'))
|
| 89 |
model.eval()
|
|
|
|
| 61 |
def RobertaClassifierOpenGPTInference(input_text):
|
| 62 |
tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
|
| 63 |
model_path = "ClassifierCheckpoint/RobertaClassifierOpenGPT.pth"
|
| 64 |
+
model = RobertaForSequenceClassification.from_pretrained('roberta-base', num_labels=2, map_location=device)
|
| 65 |
model.load_state_dict(torch.load(model_path))
|
| 66 |
model = model.to(torch.device('cpu'))
|
| 67 |
model.eval()
|
|
|
|
| 83 |
def RobertaClassifierCSAbstractInference(input_text):
|
| 84 |
tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
|
| 85 |
model_path = "ClassifierCheckpoint/RobertaClassifierCSAbstract.pth"
|
| 86 |
+
model = RobertaForSequenceClassification.from_pretrained('roberta-base', num_labels=2, map_location=device)
|
| 87 |
model.load_state_dict(torch.load(model_path))
|
| 88 |
model = model.to(torch.device('cpu'))
|
| 89 |
model.eval()
|