Upload folder using huggingface_hub

modeling_mist_finetuned.py

@@ -20,6 +20,8 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+
+AutoTokenizer.register("SmirkTokenizer", fast_tokenizer_class=SmirkTokenizerFast)
 MODEL_TYPE_ALIASES = {}
 
 def build_encoder(enc_dict: Dict[str, Any]):
@@ -121,10 +123,12 @@ class MISTFinetuned(PreTrainedModel):
         if getattr(self, "tokenizer", None) is not None:
             return self.tokenizer
         try:
-            return AutoTokenizer.from_pretrained(self.name_or_path, use_fast=True)
+            return AutoTokenizer.from_pretrained(
+                self.name_or_path, use_fast=True, trust_remote_code=True
+            )
         except Exception:
             return AutoTokenizer.from_pretrained(
-                self.config._name_or_path, use_fast=True
+                self.config._name_or_path, use_fast=True, trust_remote_code=True
             )
 
     def embed(self, smi: List[str], tokenizer=None):
@@ -142,10 +146,14 @@ class MISTFinetuned(PreTrainedModel):
     def predict(self, smi: List[str], return_dict: bool = True, tokenizer=None):
         tok = self._resolve_tokenizer(tokenizer)
         batch = tok(smi)
-        batch = DataCollatorWithPadding(tok)(batch)
-        inputs = {k: v.to(self.device) for k, v in batch.items()}
+        collate_fn = DataCollatorWithPadding(tok)
+        batch = collate_fn(batch)
+        batch = {
+            "input_ids": batch["input_ids"].to(self.encoder.device),
+            "attention_mask": batch["attention_mask"].to(self.encoder.device),
+        }
         with torch.inference_mode():
-            out = self(**inputs).cpu()
+            out = self(**batch).cpu()
         if self.channels is None or not return_dict:
             return out
         return annotate_prediction(out, maybe_get_annotated_channels(self.channels))
@@ -287,10 +295,12 @@ class MISTMultiTask(PreTrainedModel):
         if getattr(self, "tokenizer", None) is not None:
             return self.tokenizer
         try:
-            return AutoTokenizer.from_pretrained(self.name_or_path, use_fast=True)
+            return AutoTokenizer.from_pretrained(
+                self.name_or_path, use_fast=True, trust_remote_code=True
+            )
         except Exception:
             return AutoTokenizer.from_pretrained(
-                self.config._name_or_path, use_fast=True
+                self.config._name_or_path, use_fast=True, trust_remote_code=True
             )
 
     def predict(self, smi: List[str], tokenizer=None):