StarCoder2: The Versatile Code Companion.

StarCoder2 is a Transformer model designed specifically for code generation and understanding. With 13 billion parameters, it builds upon the advancements of its predecessors and is trained on a diverse dataset that includes multiple programming languages. This extensive training allows StarCoder2 to support a wide array of coding tasks, from code completion to generation.

StarCoder2 was developed to assist developers in writing and understanding code more efficiently, making it a valuable tool for various software development and data science tasks.

Pretrained models can be loaded with pretrained of the companion object:

val starcoder2 = StarCoder2Transformer.pretrained()
  .setInputCols("document")
  .setOutputCol("generation")

The default model is "StarCoder2-3B", if no name is provided. For available pretrained models please see the Models Hub.

For extended examples of usage, see StarCoder2TestSpec.

References:

• StarCoder2: The Versatile Code Companion
• https://github.com/bigcode-project/starcoder

Paper Abstract:

The BigCode project,1 an open-scientific collaboration focused on the responsible development of Large Language Models for Code (Code LLMs), introduces StarCoder2. In partnership with Software Heritage (SWH),2 we build The Stack v2 on top of the digital commons of their source code archive. Alongside the SWH repositories spanning 619 programming languages, we carefully select other high-quality data sources, such as GitHub pull requests, Kaggle notebooks, and code documentation. This results in a training set that is 4× larger than the first StarCoder dataset. We train StarCoder2 models with 3B, 7B, and 15B parameters on 3.3 to 4.3 trillion tokens and thoroughly evaluate them on a comprehensive set of Code LLM benchmarks.

We find that our small model, StarCoder2-3B, outperforms other Code LLMs of similar size on most benchmarks, and also outperforms StarCoderBase-15B. Our large model, StarCoder2- 15B, significantly outperforms other models of comparable size. In addition, it matches or outperforms CodeLlama-34B, a model more than twice its size. Although DeepSeekCoder- 33B is the best-performing model at code completion for high-resource languages, we find that StarCoder2-15B outperforms it on math and code reasoning benchmarks, as well as several low-resource languages. We make the model weights available under an OpenRAIL license and ensure full transparency regarding the training data by releasing the SoftWare Heritage persistent IDentifiers (SWHIDs) of the source code data.

Note:

This is a computationally intensive module, especially for larger code sequences. The use of an accelerator such as GPU is recommended.

Example

import spark.implicits._
import com.johnsnowlabs.nlp.base.DocumentAssembler
import com.johnsnowlabs.nlp.annotators.seq2seq.StarCoder2Transformer
import org.apache.spark.ml.Pipeline

val documentAssembler = new DocumentAssembler()
  .setInputCol("text")
  .setOutputCol("documents")

val starcoder2 = StarCoder2Transformer.pretrained("starcoder2")
  .setInputCols(Array("documents"))
  .setMinOutputLength(10)
  .setMaxOutputLength(50)
  .setDoSample(false)
  .setTopK(50)
  .setNoRepeatNgramSize(3)
  .setOutputCol("generation")

val pipeline = new Pipeline().setStages(Array(documentAssembler, starcoder2))

val data = Seq(
  "def add(a, b):"
).toDF("text")
val result = pipeline.fit(data).transform(data)

results.select("generation.result").show(truncate = false)
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
|result                                                                                                                                                                                              |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
|[def add(a, b): return a + b]                                                                                                                                                                       |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+