McGill-NLP/codellm_1b_alibi
This model is a 1B-scale decoder-only transformer designed to explore the impact of positional encoding on length generalization, specifically trained with ALiBi positional encoding to assess its effectiveness in length generalization tasks.
Usage Example
1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4model_name = "McGill-NLP/codellm_1b_alibi"
5
6# Important: `trust_remote_code=True` is required due to
7# the custom architecture supporting different positional encodings,
8# necessitating the download of the model implementation from Huggingface
9model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
10tokenizer = AutoTokenizer.from_pretrained(model_name)
11
12print(model.config.position_encoding_type)
13# Outputs: `alibi`
14
15prompt = "def print_hello_world():"
16input_ids = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).input_ids
17input_ids = torch.cat([
18 torch.tensor([[tokenizer.bos_token_id]], device="cuda"), input_ids
19], dim=1) # Prepend <bos> token
20
21output = model.generate(input_ids, do_sample=True, temperature=0.2, max_length=16)
22print(tokenizer.decode(output[0]))
Model Details
Model Description
- Developed by: McGill NLP Group
- Model type: Decoder-only transformer
- Language(s) (NLP): Primarily English, with potential application across various programming languages as demonstrated by its training on a code dataset.
- License: Apache 2.0
- Finetuned from model: This model is pretrained from scratch.
Model Sources
Uses
Direct Use
The model is designed for direct application in NLP tasks that require understanding and generating text. It's especially suited for working with source code, making it a valuable tool for tasks such as code completion, bug fixing, or even code generation.
Bias, Risks, and Limitations
Given the model's training on source code, it might inherit biases present in the underlying dataset, including but not limited to, biases towards more commonly used programming languages or coding styles. Users should be cautious when applying this model to diverse or underrepresented coding languages and contexts.
This model has not undergone safety training and it is only produced for research purposes. The user is soley responsible for outputs of this model.
Recommendations
Users should consider the context and diversity of the application domain when employing this model, especially in critical systems. Further evaluation and fine-tuning might be necessary to mitigate any potential biases or limitations for specific use cases.
How to Get Started with the Model
Use the example provided in the README to get started with generating text or code. Ensure you have the necessary dependencies installed, including torch and transformers, and follow the guidelines for setting up your environment.
Training Details
Training Data
The model was pretrained on a dataset comprising 30M source code files from the StarCoder corpus, amounting to 30B token. The training data mix:
- 40% Python
- 25% Java
- 25% JavaScript
- 5% GitHub issues
- 5% GitHub commits
Training Procedure
The model follows a decoder-only architecture with 1.3 billion parameters and was trained to predict the next token in the sequence. For more detailed information on the training procedure, refer to the paper linked above.
Technical Specifications
Model Architecture and Objective
The model leverages a decoder-only transformer architecture with ALiBi positional encoding.
Citation
Please cite the following paper if you use this model in your work:
1@inproceedings{kazemnejad2023:ImpactOfPeOnLengthGen,
2 title={The Impact of Positional Encoding on Length Generalization in Transformers},
3 author={Amirhossein Kazemnejad and Inkit Padhi and Karthikeyan Natesan and Payel Das and Siva Reddy},
4 booktitle={Thirty-seventh Conference on Neural Information Processing Systems},
5 year={2023},
6 url={https://openreview.net/forum?id=Drrl2gcjzl}
7}
More Information
For further details about the model's architecture, training, and applications, please refer to the paper and the GitHub repository linked above.
