Update README.md

README.md

@@ -130,23 +130,33 @@ export TOKENIZERS_PARALLELISM=false
 ### 2. Load the Model via Hugging Face
 
 ```python
+import os
+import warnings
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+from transformers import logging as hf_logging
+hf_logging.set_verbosity_error() 
+
+warnings.filterwarnings("ignore", category=UserWarning)
 import torch
 from transformers import AutoTokenizer
 from proto_model.configuration_proto import ProtoConfig
 from proto_model.modeling_proto    import ProtoForMultiLabelClassification
 
-cfg       = ProtoConfig.from_pretrained("row56/ProtoPatient")
+cfg = ProtoConfig.from_pretrained("row56/ProtoPatient")
 cfg.pretrained_model_name_or_path = "bert-base-uncased"
-cfg.use_cuda                     = False
+cfg.use_cuda = torch.cuda.is_available()
+
+device = torch.device("cuda" if cfg.use_cuda else "cpu")
 
 tokenizer = AutoTokenizer.from_pretrained(cfg.pretrained_model_name_or_path)
-model     = ProtoForMultiLabelClassification.from_pretrained(
-                "row56/ProtoPatient",
-                config=cfg,
-                ignore_mismatched_sizes=True
-            )
+model = ProtoForMultiLabelClassification.from_pretrained(
+    "row56/ProtoPatient",
+    config=cfg,
+)
+model.to(device)
 model.eval()
-model.cpu()
 
 def get_proto_logits(texts):
     enc = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
@@ -167,9 +177,6 @@ texts = [
 logits = get_proto_logits(texts)
 print("Logits shape:", logits.shape)
 print("Logits:\n", logits)
-
-probs = torch.sigmoid(logits)
-print("Probabilities:\n", probs)
 ```
 
 ## 3. Training Data & Licenses