import unittest
from deepchem.feat.roberta_tokenizer import RobertaFeaturizer
import pytest


class TestRobertaFeaturizer(unittest.TestCase):
  """Tests for RobertaFeaturizer"""

  @pytest.mark.torch
  def setUp(self):
    self.smiles = ["Cn1c(=O)c2c(ncn2C)n(C)c1=O", "CC(=O)N1CN(C(C)=O)C(O)C1O"]
    self.long_molecule_smiles = [
        "CCCCCCCCCCCCCCCCCCCC(=O)OCCCNC(=O)c1ccccc1SSc1ccccc1C(=O)NCCCOC(=O)CCCCCCCCCCCCCCCCCCC"
    ]
    self.featurizer = RobertaFeaturizer.from_pretrained(
        "seyonec/SMILES_tokenized_PubChem_shard00_160k")

  @pytest.mark.torch
  def test_smiles_call(self):
    """Test __call__ method for the featurizer, which is inherited from HuggingFace's RobertaTokenizerFast"""
    embedding = self.featurizer(
        self.smiles, add_special_tokens=True, truncation=True)
    embedding_long = self.featurizer(
        self.long_molecule_smiles * 2, add_special_tokens=True, truncation=True)
    for emb in [embedding, embedding_long]:
      assert 'input_ids' in emb.keys() and 'attention_mask' in emb.keys()
      assert len(emb['input_ids']) == 2 and len(emb['attention_mask']) == 2

  @pytest.mark.torch
  def test_smiles_featurize(self):
    """Test the .featurize method, which will convert the dictionary output to an array

    Checks that all SMILES are featurized and that each featurization
    contains input_ids and attention_mask
    """
    feats = self.featurizer.featurize(
        self.smiles, add_special_tokens=True, truncation=True)
    assert (len(feats) == 2)
    assert (all([len(f) == 2 for f in feats]))

    long_feat = self.featurizer.featurize(
        self.long_molecule_smiles, add_special_tokens=True, truncation=True)
    assert (len(long_feat) == 1)
    assert (len(long_feat[0] == 2))

@@ -4,17 +4,57 @@
 ort 
-unittest%0A
+pytest%0A%0A%0A@pytest.mark.torch%0Adef setUp(self):%0A    
 from
@@ -115,149 +115,8 @@
 zer%0A
-import pytest%0A%0A%0Aclass TestRobertaFeaturizer(unittest.TestCase):%0A  %22%22%22Tests for RobertaFeaturizer%22%22%22%0A%0A  @pytest.mark.torch%0A  def setUp(self):%0A
     
@@ -441,18 +441,17 @@
 160k%22)%0A%0A
-  
+%0A
 @pytest.
@@ -453,34 +453,32 @@
 test.mark.torch%0A
-  
 def test_smiles_
@@ -595,24 +595,419 @@
 izerFast%22%22%22%0A
+    from deepchem.feat.roberta_tokenizer import RobertaFeaturizer%0A    self.smiles = %5B%22Cn1c(=O)c2c(ncn2C)n(C)c1=O%22, %22CC(=O)N1CN(C(C)=O)C(O)C1O%22%5D%0A    self.long_molecule_smiles = %5B%0A        %22CCCCCCCCCCCCCCCCCCCC(=O)OCCCNC(=O)c1ccccc1SSc1ccccc1C(=O)NCCCOC(=O)CCCCCCCCCCCCCCCCCCC%22%0A    %5D%0A    self.featurizer = RobertaFeaturizer.from_pretrained(%0A        %22seyonec/SMILES_tokenized_PubChem_shard00_160k%22)%0A
     embeddin
@@ -1255,24 +1255,26 @@
 ong%5D:%0A      
+  
 assert 'inpu
@@ -1329,16 +1329,18 @@
 .keys()%0A
+  
       as
@@ -1408,18 +1408,17 @@
 ) == 2%0A%0A
+%0A
-  
 @pytest.
@@ -1428,18 +1428,16 @@
 k.torch%0A
-  
 def test
@@ -1666,24 +1666,419 @@
 ask%0A    %22%22%22%0A
+    from deepchem.feat.roberta_tokenizer import RobertaFeaturizer%0A    self.smiles = %5B%22Cn1c(=O)c2c(ncn2C)n(C)c1=O%22, %22CC(=O)N1CN(C(C)=O)C(O)C1O%22%5D%0A    self.long_molecule_smiles = %5B%0A        %22CCCCCCCCCCCCCCCCCCCC(=O)OCCCNC(=O)c1ccccc1SSc1ccccc1C(=O)NCCCOC(=O)CCCCCCCCCCCCCCCCCCC%22%0A    %5D%0A    self.featurizer = RobertaFeaturizer.from_pretrained(%0A        %22seyonec/SMILES_tokenized_PubChem_shard00_160k%22)%0A
     feats = 
@@ -2243,17 +2243,16 @@
 eats%5D))%0A
-%0A
     long

change from unittest to pytest