NT Java 1.1B

🚀 NT-Java-1.1B

NT-Java-1.1Bは、Javaプログラミングのコーディングタスク向けに設計されたオープンソースの専用コードモデルです。StarcoderBase-1Bを拡張して事前学習を行い、Javaのコーディングに特化した性能を備えています。

🚀 クイックスタート

NT-Java-1.1Bは、Javaのコーディングタスクに特化したオープンソースのモデルです。このモデルは、StarcoderBase-1Bを拡張して事前学習を行っており、Javaのコード生成や補完に優れた性能を発揮します。

✨ 主な機能

• Java特化: Javaプログラミングタスクに特化した性能を持ち、Javaのコード生成や補完に適しています。
• デスクトップデプロイ可能: 小型言語モデル（SLM）であるため、消費者向けPCにデプロイできます。
• 高い性能: 同等サイズのオープンソースコードモデルを上回る性能を示します。

📦 インストール

以下のコマンドを使用して必要なライブラリをインストールします。

# pip install -q transformers

💻 使用例

基本的な使用法

コード生成

# pip install -q transformers
from transformers import AutoModelForCausalLM, AutoTokenizer

checkpoint = "infosys/NT-Java-1.1B"
device = "cuda" # for GPU usage or "cpu" for CPU usage

tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)

inputs = tokenizer.encode("public class HelloWorld {\n    public static void main(String[] args) {", return_tensors="pt").to(device)
outputs = model.generate(inputs)
print(tokenizer.decode(outputs[0]))

中間埋め込み（Fill-in-the-middle）

input_text = "<fim_prefix>public class PalindromeChecker {\n        public static boolean isPalindrome(String str) {\n          <fim_suffix>return true;\n      }\n<fim_middle>"
inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
outputs = model.generate(inputs)
print(tokenizer.decode(outputs[0]))

高度な使用法

8ビット量子化バージョンの使用

# pip install bitsandbytes accelerate
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

# to use 4bit use `load_in_4bit=True` instead
quantization_config = BitsAndBytesConfig(load_in_8bit=True)

checkpoint = "infosys/NT-Java-1.1B"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForCausalLM.from_pretrained(checkpoint, quantization_config=quantization_config)

inputs = tokenizer.encode("public class HelloWorld {\n    public static void main(String[] args) {", return_tensors="pt").to("cuda")
outputs = model.generate(inputs)
print(tokenizer.decode(outputs[0]))

📚 ドキュメント

モデル情報

モデルの意図された用途

大規模コードモデルは推論にGPUなどの特殊なハードウェアを必要としますが、NT-Java-1.1Bは小型言語モデルであるため、消費者向けPCにデプロイできます。また、Javaプログラミングタスクにおいて、同等サイズのオープンソースコードモデルを上回る性能を示します。

モデルの訓練

モデル

• アーキテクチャ: Multi-Query AttentionとFill-in-the-Middle目的を持つGPT-2モデル。
• 訓練ステップ: 120K
• コンテキスト長: 8Kトークン
• 事前学習トークン: 220億
• 精度: bfloat16

ハードウェア

• GPU: 6台のNVIDIA A100 80GB
• 訓練時間: 10日

ソフトウェア

• オーケストレーション: Megatron-LM
• ニューラルネットワーク: PyTorch

帰属とその他の要件

モデルの事前学習データセットは、許容的なライセンスのデータのみを含むように選別されています。ただし、モデルはデータセットからのソースコードをそのまま生成する可能性があります。そのようなコードのライセンスには、帰属や他の特定の条件が必要になる場合があります。これに対応するため、BigCodeは検索インデックスを提供しており、ユーザーは生成されたコードの事前学習データ内での起源を追跡できます。

モデルの制限

NT-Java-1.1Bモデルは公開されているデータセットで訓練されており、安全保証は一切ありません。すべての言語モデルと同様に、その出力は本質的に予測不可能であり、生成されたコードは期待どおりに動作しない可能性があります。また、コードは非効率的であったり、バグやセキュリティホールを含んでいる可能性があります。したがって、ユーザーや開発者は十分な安全テストを行い、特定のニーズに合った強力なフィルタリングメカニズムを実装することが重要です。

🔧 技術詳細

NT-Java-1.1Bは、StarcoderBase-1Bを拡張して事前学習を行ったモデルです。モデルは、Multi-Query AttentionとFill-in-the-Middle目的を持つGPT-2アーキテクチャを使用しています。また、JavaサブセットのStarcoderDataデータセットで訓練されています。

📄 ライセンス

このモデルは、BigCode OpenRAIL-M v1ライセンス契約の下で提供されています。完全な契約はこちらで確認できます。

📖 引用

@article{rathinasamy2024narrow,
  title={Narrow Transformer: Starcoder-Based Java-LM For Desktop},
  author={Kamalkumar Rathinasamy and Balaji A J and Rajab Ali Mondal and Ankush Kumar and Harshini K and Gagan Gayari and Sreenivasa Raghavan Karumboor Seshadri and Swayam Singh},
  journal={arXiv preprint arXiv:2407.03941},
  year={2024}
}