Fixes test on ci

deepchem/feat/reaction_featurizer.py

@@ -1,6 +1,7 @@
 from deepchem.feat import Featurizer
 from typing import List
 import numpy as np
+from numpy.typing import ArrayLike
 
 try:
     from transformers import RobertaTokenizerFast
@@ -35,7 +36,10 @@ class RxnFeaturizer(Featurizer):
         - False - Mix the reactants and reagents
     """
 
-    def __init__(self, tokenizer: RobertaTokenizerFast, sep_reagent: bool):
+    def __init__(self,
+                 tokenizer: RobertaTokenizerFast,
+                 sep_reagent: bool,
+                 max_length: int = 100):
         """Initialize a ReactionFeaturizer object.
 
         Parameters
@@ -44,6 +48,8 @@ def __init__(self, tokenizer: RobertaTokenizerFast, sep_reagent: bool):
             HuggingFace Tokenizer to be used for featurization.
         sep_reagent: bool
             Toggle to separate or mix the reactants and reagents.
+        max_length: int, default 100
+            Maximum length of padding
         """
         if not isinstance(tokenizer, RobertaTokenizerFast):
             raise TypeError(
@@ -52,8 +58,9 @@ def __init__(self, tokenizer: RobertaTokenizerFast, sep_reagent: bool):
         else:
             self.tokenizer = tokenizer
         self.sep_reagent = sep_reagent
+        self.max_length = max_length
 
-    def _featurize(self, datapoint: str, **kwargs) -> List[List[List[int]]]:
+    def _featurize(self, datapoint: str, **kwargs) -> List[ArrayLike]:
         """Featurizes a datapoint.
 
         Processes each entry in the dataset by first applying the reactant-reagent
@@ -87,15 +94,26 @@ def _featurize(self, datapoint: str, **kwargs) -> List[List[List[int]]]:
             ]
         target = product
 
-        source_encoding = list(
-            self.tokenizer(source, padding=True, **kwargs).values())
-        target_encoding = list(
-            self.tokenizer(target, padding=True, **kwargs).values())
-
+        source_encoding = np.asarray(
+            list(
+                self.tokenizer(source,
+                               padding='max_length',
+                               truncation=True,
+                               max_length=self.max_length,
+                               **kwargs).values()))
+        target_encoding = np.asarray(
+            list(
+                self.tokenizer(target,
+                               padding='max_length',
+                               truncation=True,
+                               max_length=self.max_length,
+                               **kwargs).values()))
         return [source_encoding, target_encoding]
 
     def __call__(self, *args, **kwargs) -> np.ndarray:
-        return self.featurize(*args, **kwargs)
+        features = self.featurize(*args, **kwargs)
+        print(type(features), len(features))
+        return features
 
     def __str__(self) -> str:
         """Handles file name error.

deepchem/feat/tests/test_reaction_featurizer.py

@@ -10,10 +10,13 @@ def test_featurize():
     from deepchem.feat.reaction_featurizer import RxnFeaturizer
     tokenizer = RobertaTokenizerFast.from_pretrained(
         "seyonec/PubChem10M_SMILES_BPE_450k")
-    featurizer = RxnFeaturizer(tokenizer, sep_reagent=True)
+    max_length = 20
+    featurizer = RxnFeaturizer(tokenizer,
+                               sep_reagent=True,
+                               max_length=max_length)
     reaction = ['CCS(=O)(=O)Cl.OCCBr>CCN(CC)CC.CCOCC>CCS(=O)(=O)OCCBr']
     feats = featurizer.featurize(reaction)
-    assert (feats.shape == (1, 2, 2, 1))
+    assert (feats.shape == (1, 2, 2, 1, max_length))
 
 
 @pytest.mark.torch
@@ -33,7 +36,7 @@ def test_separation():
     feats_sep = featurizer_sep.featurize(reaction)
 
     # decode the source in the mixed and separated cases
-    mix_decoded = tokenizer.decode(feats_mix[0][0][0][0])
-    sep_decoded = tokenizer.decode(feats_sep[0][0][0][0])
+    mix_decoded = tokenizer.decode(feats_mix[0][0][0][0]).replace('<pad>', '')
+    sep_decoded = tokenizer.decode(feats_sep[0][0][0][0]).replace('<pad>', '')
     assert mix_decoded == '<s>CCS(=O)(=O)Cl.OCCBr.CCN(CC)CC.CCOCC></s>'
     assert sep_decoded == '<s>CCS(=O)(=O)Cl.OCCBr>CCN(CC)CC.CCOCC</s>'

deepchem/feat/tests/test_roberta_tokenizer.py

@@ -34,13 +34,20 @@ def test_smiles_featurize():
     ]
     featurizer = RobertaFeaturizer.from_pretrained(
         "seyonec/SMILES_tokenized_PubChem_shard00_160k")
-    feats = featurizer.featurize(smiles,
-                                 add_special_tokens=True,
-                                 truncation=True)
-    assert (len(feats) == 2)
-    assert (all([len(f) == 2 for f in feats]))
-    long_feat = featurizer.featurize(long_molecule_smiles,
-                                     add_special_tokens=True,
-                                     truncation=True)
-    assert (len(long_feat) == 1)
-    assert (len(long_feat[0] == 2))
+    max_length = 100
+    feat_kwargs = {
+        'add_special_tokens': True,
+        'truncation': True,
+        'padding': 'max_length',
+        'max_length': max_length
+    }
+    feats = featurizer.featurize(smiles, **feat_kwargs)
+    assert len(feats) == 2
+    assert all([len(f) == 2 for f in feats])
+    assert all([len(f[0]) == max_length for f in feats])
+
+    long_feat = featurizer.featurize(long_molecule_smiles, **feat_kwargs)
+    assert len(long_feat) == 1
+    assert len(long_feat[0]) == 2  # the tokens and attention mask
+    assert len(
+        long_feat[0][0]) == 100  # number of tokens for each smiles string

deepchem/models/gbdt_models/gbdt_model.py

@@ -115,7 +115,7 @@ def fit(self, dataset: Dataset):
 
         # retrain model to whole data using best n_estimators * 1.25
         if self.model.__class__.__name__.startswith('XGB'):
-            estimated_best_round = np.round(self.model.best_ntree_limit * 1.25)
+            estimated_best_round = np.round(self.model.best_iteration * 1.25)
         else:
             estimated_best_round = np.round(self.model.best_iteration_ * 1.25)
         self.model.n_estimators = np.int64(estimated_best_round)