Cannot load large XlmRoBertaForTokenClassification model in scala 2.12

Steps to Reproduce

1. XlmRoBertaForTokenClassification.loadSavedModel(<large model path>)

Stack Trace

Exception in thread “main” java.lang.OutOfMemoryError: Java heap space at com.johnsnowlabs.ml.tensorflow.io.ChunkBytes$.readFileInByteChunks(ChunkBytes.scala:44) at com.johnsnowlabs.ml.tensorflow.TensorflowWrapper$.read(TensorflowWrapper.scala:436) at com.johnsnowlabs.nlp.annotators.classifier.dl.ReadXlmRoBertaForTokenTensorflowModel.loadSavedModel(XlmRoBertaForTokenClassification.scala:311) at com.johnsnowlabs.nlp.annotators.classifier.dl.ReadXlmRoBertaForTokenTensorflowModel.loadSavedModel$(XlmRoBertaForTokenClassification.scala:292) at com.johnsnowlabs.nlp.annotators.classifier.dl.XlmRoBertaForTokenClassification$.loadSavedModel(XlmRoBertaForTokenClassification.scala:330)

Your Environment

• Spark NLP version 3.3.1:
• Apache NLP version 3.0.1:
• Java version 1.8.0:
• Setup and installation (Pypi, Conda, Maven, etc.): SBT + Scala
• Operating System and version: Ubuntu + MacOS

[Pcosmin]

You are write. I added -Xmx15g to the java process and it’s working. Thank you!

[maziyarpanahi]

Thanks, the error clearly indicates there is not enough memory to serialize XLM-RoBERTa large model in Java. It seems strange, the 15G should be enough for that model. Might be some settings regarding Java heap in your classpath or you actually don’t have 15G actual free memory.

I just tested xlm-roberta-large on Google Colab which only has 12G memory and 2G-3G was already in use by previous operations and it worked: https://colab.research.google.com/drive/1p5jFqxMuCnfcWFDGy_JS7yLJDeYGF00f?usp=sharing

Unfortunately, not much left other than freeing up more memory on that machine.