NT Java 1.1B

🚀 NT-Java-1.1B 代碼模型

NT-Java-1.1B 是一款專為 Java 編程任務打造的開源代碼模型。它基於 StarCoderBase-1B 進行擴展預訓練，在 Java 編碼方面表現出色，且能部署在消費級 PC 上，為開發者提供便利。

🚀 快速開始

NT-Java-1.1B 是基於 StarCoderBase-1B 擴展預訓練得到的開源代碼模型，專為 Java 編程任務設計。該模型是僅解碼器的 Transformer，採用了 Multi-Query Attention，上下文長度達 8192 個標記。它使用 StarCoderData 數據集中的 Java 子集進行訓練，約有 220 億個標記。

✨ 主要特性

• 可部署性強：作為小型語言模型（SLM），NT-Java-1.1B 能夠部署在消費級 PC 上，無需依賴如 GPU 等專業硬件進行推理。
• 性能出色：在 Java 編程任務中，該模型的表現優於同等規模的開源代碼模型。
• 多功能用途：適用於內存/計算資源受限的環境、對延遲敏感的場景，以及 Java 代碼生成、補全和代碼填充（FIM）等任務。
• 量化版本靈活：量化版本 NT-Java-1.1B-GGUF 在 MultiPL-E Java 代碼基準測試中表現與 10 億參數的開源模型相當，且可與 Ollama、GPT4ALL 等多個框架配合使用。

📦 安裝指南

使用該模型前，需安裝必要的庫：

pip install -q transformers

若使用量化版本，還需安裝：

pip install bitsandbytes accelerate

💻 使用示例

基礎用法

# 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）

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]))

🔧 技術細節

模型

• 架構：採用帶有 Multi-Query Attention 和 Fill-in-the-Middle 目標的 GPT - 2 模型。
• 訓練步驟：120K
• 上下文長度：8K 標記
• 預訓練標記：220 億
• 精度：bfloat16

硬件

• GPU：6 塊 NVIDIA A100 80GB
• 訓練時間：10 天

軟件

• 編排工具：Megatron-LM
• 神經網絡框架：PyTorch

📄 許可證

該模型遵循 BigCode OpenRAIL - M v1 許可協議，你可以在 此處 查看完整協議。

⚠️ 重要提示

模型的預訓練數據集僅包含具有寬鬆許可的數據，但模型仍可能逐字生成數據集中的源代碼。此類代碼的許可證可能要求進行歸因並遵守其他特定條件。為方便合規，BigCode 提供了一個 搜索索引，用戶可通過該索引追蹤生成代碼在預訓練數據中的來源，以進行正確歸因並遵守許可要求。

💡 使用建議

NT-Java-1.1B 模型是在公開可用的數據集上訓練的，不提供任何安全保證。由於所有語言模型的輸出都具有不可預測性，生成的代碼可能無法按預期執行，還可能存在效率低下、包含錯誤和安全漏洞等問題。因此，用戶和開發者必須進行全面的安全測試，並根據自身需求實施強大的過濾機制。

📖 引用

@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}
}