Add device

README.md

@@ -53,8 +53,9 @@ class WBGDocTopic:
             each document, along with the mean and standard deviation of the topic classification scores.
     """
 
-    def __init__(self, classifiers: dict = None):
+    def __init__(self, classifiers: dict = None, device: str = None
         self.classifiers = classifiers or {}
+        self.device = device
 
         if classifiers is None:
             self.load_classifiers()
@@ -71,7 +72,7 @@ class WBGDocTopic:
                     continue
 
                 model_name = f"avsolatorio/doc-topic-model_eval-{i:02}_train-{j:02}"
-                classifier = pipeline("text-classification", model=model_name, tokenizer=tokenizer, top_k=None)
+                classifier = pipeline("text-classification", model=model_name, tokenizer=tokenizer, top_k=None, device=self.device)
 
                 self.classifiers[model_name] = classifier
 
@@ -118,7 +119,8 @@ sample_text = """A growing literature attributes gender inequality in labor mark
 sents = sent_tokenize(inp)
 
 # Create the instance which will load the models.
-dtopic_model = WBGDocTopic()
+# Set the device to "cuda" if you want to use a GPU.
+dtopic_model = WBGDocTopic(device=None)
 
 # Infer the topics and scores
 outs = dtopic_model.suggest_topics(sents)