jjzha commited on
Commit
4413a8c
1 Parent(s): 56a2dbb

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -2
app.py CHANGED
@@ -1,8 +1,11 @@
1
  import gradio as gr
 
2
  from transformers import pipeline
3
 
4
- token_skill_classifier = pipeline(model="jjzha/escoxlmr_skill_extraction", aggregation_strategy="first", device="cuda")
5
- token_knowledge_classifier = pipeline(model="jjzha/escoxlmr_knowledge_extraction", aggregation_strategy="first", device="cuda")
 
 
6
 
7
 
8
  examples = [
@@ -32,6 +35,7 @@ def aggregate_span(results):
32
  return new_results
33
 
34
 
 
35
  def ner(text):
36
  output_skills = token_skill_classifier(text)
37
  for result in output_skills:
 
1
  import gradio as gr
2
+ import spaces
3
  from transformers import pipeline
4
 
5
+ token_skill_classifier = pipeline(model="jjzha/escoxlmr_skill_extraction", aggregation_strategy="first")
6
+ token_knowledge_classifier = pipeline(model="jjzha/escoxlmr_knowledge_extraction", aggregation_strategy="first")
7
+ token_skill_classifier.to("cuda")
8
+ token_knowledge_classifier.to("cuda")
9
 
10
 
11
  examples = [
 
35
  return new_results
36
 
37
 
38
+ @spaces.GPU
39
  def ner(text):
40
  output_skills = token_skill_classifier(text)
41
  for result in output_skills: