Spaces:

devngho
/

ko-edu-classifier

Sleeping

devngho commited on Sep 29

Commit

baaa2f1

•

1 Parent(s): 2bde448

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -1,14 +1,17 @@
 from transformers import pipeline
 import gradio as gr
 models = {
-    'devngho/ko_edu_classifier_v2_nlpai-lab_KoE5': pipeline("text-classification", model="devngho/ko_edu_classifier_v2_nlpai-lab_KoE5"),
-    'devngho/ko_edu_classifier_v2_lemon-mint_LaBSE-EnKo-Nano-Preview-v0.3': pipeline("text-classification", model="devngho/ko_edu_classifier_v2_lemon-mint_LaBSE-EnKo-Nano-Preview-v0.3"),
-    'devngho/ko_edu_classifier_v2_LaBSE': pipeline("text-classification", model="devngho/ko_edu_classifier_v2_LaBSE")
 }
 import gradio as gr
 def evaluate_model(input_text):
     return [model(input_text)[0]['score'] * 6 if model_name != 'devngho/ko_edu_classifier_v2_nlpai-lab_KoE5' else model('passage: ' + input_text)[0]['score'] * 6 for model_name, model in models.items()]

 from transformers import pipeline
 import gradio as gr
+import spaces
+import torch
 models = {
+    'devngho/ko_edu_classifier_v2_nlpai-lab_KoE5': pipeline("text-classification", model="devngho/ko_edu_classifier_v2_nlpai-lab_KoE5", device='cuda', torch_dtype=torch.bfloat16),
+    'devngho/ko_edu_classifier_v2_lemon-mint_LaBSE-EnKo-Nano-Preview-v0.3': pipeline("text-classification", model="devngho/ko_edu_classifier_v2_lemon-mint_LaBSE-EnKo-Nano-Preview-v0.3", device='cuda', torch_dtype=torch.bfloat16),
+    'devngho/ko_edu_classifier_v2_LaBSE': pipeline("text-classification", model="devngho/ko_edu_classifier_v2_LaBSE", device='cuda', torch_dtype=torch.bfloat16)
 }
 import gradio as gr
+@spaces.GPU
 def evaluate_model(input_text):
     return [model(input_text)[0]['score'] * 6 if model_name != 'devngho/ko_edu_classifier_v2_nlpai-lab_KoE5' else model('passage: ' + input_text)[0]['score'] * 6 for model_name, model in models.items()]