Update versions

app.py

@@ -1,6 +1,7 @@
 from logging import PlaceHolder
 import gradio as gr
 import os, sys
+import spaces
 from npc_bert_models.gradio_demo import *
 from npc_bert_models.mlm_module import NpcBertMLM
 from npc_bert_models.cls_module import NpcBertCLS
@@ -9,7 +10,7 @@ from npc_bert_models.app_logger import get_logger
 import json
 
 class main_window():
-    logger = get_logger('main')
+    logger = get_logger('main', log_level='debug')
     def __init__(self):
         self.interface = None
         self.examples = json.load(open("examples.json", 'r'))

npc_bert_models/cls_module.py

@@ -50,6 +50,7 @@ class NpcBertCLS():
         self.tokenizer = AutoTokenizer.from_pretrained(self.pretrained_model)
         try:
             self.pipeline = hf_pipeline("text-classification", model=self.model, tokenizer=self.tokenizer, device='cuda')
+            self.pipeline.model.to('cuda')
         except Exception as e:
             self.pipeline = hf_pipeline("text-classification", model=self.model, tokenizer=self.tokenizer, device='cpu')
             self.logger.warning("No GPU!")

npc_bert_models/mlm_module.py

@@ -49,6 +49,7 @@ class NpcBertMLM():
         self.tokenizer = AutoTokenizer.from_pretrained(self.pretrained_model)
         try:
             self.pipeline = hf_pipeline("fill-mask", model=self.model, tokenizer=self.tokenizer, device='cuda')
+            self.pipeline.model.to('cuda')
         except Exception as e:
             self.pipeline = hf_pipeline("fill-mask", model=self.model, tokenizer=self.tokenizer, device='cpu')
             self.logger.warning("No GPU")
@@ -75,6 +76,7 @@ class NpcBertMLM():
         if self.pipeline is None:
             msg = "Model was not initialized, have you run load()?"
             raise BrokenPipeError(msg)
+        
         pipe_out = self.pipeline(*args)
         # Just use the first output
         if not isinstance(pipe_out[0], dict):

npc_bert_models/summary_module.py

@@ -43,6 +43,7 @@ class NpcBertGPT2():
                                         early_stopping=True,
                                         no_repeat_ngram_size=5, 
                                         max_new_tokens=60)
+            self.pipeline.model.to('cuda')
         except Exception as e:
             self.pipeline = hf_pipeline("text2text-generation", 
                                     model=self.model, 
@@ -80,7 +81,7 @@ class NpcBertGPT2():
             raise BrokenPipeError(msg)
         
         self.logger.info(f"Called with arguments {args = }")
-        
+        self.logger.info("Model: {self.pipeline.model}")
         pipe_out, = self.pipeline(*args)
         pipe_out = pipe_out['generated_text']
         self.logger.info(f"Generated text: {pipe_out}")

requirements.txt

@@ -1,8 +1,7 @@
-torch >= 1.12.1
+torch == 2.2.2
 scikit-learn >= 1.4.0
 pandas >= 2.1.4
-transformers >= 4.37.2
-numpy >= 1.26
+transformers >= 4.37.2, < 4.50, < 2.0
 gradio >= 4.18, < 4.50
 scipy >= 1.12
 spaces