Update README.md

README.md

@@ -79,8 +79,10 @@ from torch import nn
 from torch.nn import CrossEntropyLoss
 from transformers import AutoConfig, AutoTokenizer
 from transformers.file_utils import ModelOutput
-from transformers.models.xlnet.modeling_xlnet import XLNetModel, XLNetPreTrainedModel
-
+from transformers.models.roberta.modeling_roberta import (
+    RobertaModel,
+    RobertaPreTrainedModel,
+)
 
 @dataclass
 class XGECToROutput(ModelOutput):
@@ -203,7 +205,7 @@ class XGECToRRoberta(RobertaPreTrainedModel):
 
 config = AutoConfig.from_pretrained("manred1997/roberta-large_lemon-spell_5k")
 tokenizer = AutoTokenizer.from_pretrained("manred1997/roberta-large_lemon-spell_5k")
-model = XGECToRXLNet.from_pretrained(
+model = XGECToRRoberta.from_pretrained(
     "manred1997/roberta-large_lemon-spell_5k", config=config
 )
 ```