Upload model
Browse files- MoEModel.py +4 -1
MoEModel.py
CHANGED
|
@@ -73,8 +73,11 @@ class MoeModel(PreTrainedModel):
|
|
| 73 |
return ret
|
| 74 |
|
| 75 |
|
|
|
|
|
|
|
|
|
|
| 76 |
def perplexity(model, tokenized_input) -> torch.Tensor:
|
| 77 |
with torch.inference_mode():
|
| 78 |
-
output = model(tokenized_input, labels=tokenized_input)
|
| 79 |
ppl = torch.exp(output.loss)
|
| 80 |
return ppl.item()
|
|
|
|
| 73 |
return ret
|
| 74 |
|
| 75 |
|
| 76 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 77 |
+
|
| 78 |
+
|
| 79 |
def perplexity(model, tokenized_input) -> torch.Tensor:
|
| 80 |
with torch.inference_mode():
|
| 81 |
+
output = model(tokenized_input.to(device), labels=tokenized_input)
|
| 82 |
ppl = torch.exp(output.loss)
|
| 83 |
return ppl.item()
|