Spaces:

esun-choi
/

InvoHide_Inisw_teem8

Sleeping

esun-choi commited on Dec 19, 2023

Commit

90881fe

•

1 Parent(s): 4cfe9f7

Update ner.py

Files changed (1) hide show

ner.py CHANGED Viewed

@@ -2,7 +2,8 @@ from transformers import AutoTokenizer, AutoModelForTokenClassification
 from transformers import pipeline
 from collections import defaultdict
 import torch
-device = torch.device("cuda")
 tokenizer = AutoTokenizer.from_pretrained("Leo97/KoELECTRA-small-v3-modu-ner")
 model = AutoModelForTokenClassification.from_pretrained("Leo97/KoELECTRA-small-v3-modu-ner")
 model.to(device)

 from transformers import pipeline
 from collections import defaultdict
 import torch
+# device = torch.device("cuda")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 tokenizer = AutoTokenizer.from_pretrained("Leo97/KoELECTRA-small-v3-modu-ner")
 model = AutoModelForTokenClassification.from_pretrained("Leo97/KoELECTRA-small-v3-modu-ner")
 model.to(device)