Compatibility to ZERO

npc_bert_models/cls_module.py

@@ -1,7 +1,9 @@
+import spaces.zero
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 from transformers import pipeline as hf_pipeline
 from pathlib import Path
 from typing import Any, Dict
+import spaces
 from .app_logger import get_logger
 
 class NpcBertCLS():
@@ -46,8 +48,14 @@ class NpcBertCLS():
         
         self.model = AutoModelForSequenceClassification.from_pretrained(self.pretrained_model)
         self.tokenizer = AutoTokenizer.from_pretrained(self.pretrained_model)
-        self.pipeline = hf_pipeline("text-classification", model=self.model, tokenizer=self.tokenizer, device='cpu')
+        try:
+            self.pipeline = hf_pipeline("text-classification", model=self.model, tokenizer=self.tokenizer, device='cuda')
+        except Exception as e:
+            self.pipeline = hf_pipeline("text-classification", model=self.model, tokenizer=self.tokenizer, device='cpu')
+            self.logger.warning("No GPU!")
+            self.logger.exception(e)
     
+    @spaces.GPU
     def __call__(self, *args: Any) -> Any:
         """Performs classification on the given reports.
 

npc_bert_models/mlm_module.py

@@ -2,6 +2,7 @@ from transformers import AutoTokenizer, AutoModelForMaskedLM
 from transformers import pipeline as hf_pipeline
 from pathlib import Path
 from .app_logger import get_logger
+import spaces
 
 class NpcBertMLM():
     r"""A class for performing masked language modeling with BERT.
@@ -46,8 +47,14 @@ class NpcBertMLM():
         
         self.model = AutoModelForMaskedLM.from_pretrained(self.pretrained_model)
         self.tokenizer = AutoTokenizer.from_pretrained(self.pretrained_model)
-        self.pipeline = hf_pipeline("fill-mask", model=self.model, tokenizer=self.tokenizer, device='cpu')
+        try:
+            self.pipeline = hf_pipeline("fill-mask", model=self.model, tokenizer=self.tokenizer, device='cuda')
+        except Exception as e:
+            self.pipeline = hf_pipeline("fill-mask", model=self.model, tokenizer=self.tokenizer, device='cpu')
+            self.logger.warning("No GPU")
+            self.logger.exception(e)
     
+    @spaces.GPU
     def __call__(self, *args):
         """Performs masked language modeling prediction.
 

npc_bert_models/summary_module.py

@@ -1,6 +1,7 @@
 from transformers import AutoTokenizer, EncoderDecoderModel
 from transformers import pipeline as hf_pipeline
 from pathlib import Path
+import spaces
 import re
 from .app_logger import get_logger
 
@@ -29,7 +30,21 @@ class NpcBertGPT2():
         
         self.model = EncoderDecoderModel.from_pretrained(self.pretrained_model)
         self.tokenizer = AutoTokenizer.from_pretrained(self.pretrained_model)
-        self.pipeline = hf_pipeline("text2text-generation", 
+        
+        try:
+            self.pipeline = hf_pipeline("text2text-generation", 
+                                        model=self.model, 
+                                        tokenizer=self.tokenizer, 
+                                        device='cuda', 
+                                        num_beams=4,
+                                        do_sample=True,
+                                        top_k = 5,
+                                        temperature=.95,
+                                        early_stopping=True,
+                                        no_repeat_ngram_size=5, 
+                                        max_new_tokens=60)
+        except Exception as e:
+            self.pipeline = hf_pipeline("text2text-generation", 
                                     model=self.model, 
                                     tokenizer=self.tokenizer, 
                                     device='cpu', 
@@ -40,7 +55,10 @@ class NpcBertGPT2():
                                     early_stopping=True,
                                     no_repeat_ngram_size=5, 
                                     max_new_tokens=60)
+            self.logger.warning("No GPU!")
+            self.logger.exception(e)
     
+    @spaces.GPU
     def __call__(self, *args):
         """Performs masked language modeling prediction.
 

requirements.txt

@@ -4,4 +4,6 @@ pandas >= 2.1.4
 transformers >= 4.37.2
 numpy >= 1.26
 gradio >= 4.18, < 4.50
-scipy >= 1.12
+scipy >= 1.12
+spaces
+python == 3.10.13