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