This is a repository for the CodeLlama-7b model fine-tuned on the KStack-clean dataset with rule-based filtering, in the Hugging Face Transformers format. KStack-clean is a small subset of KStack, the largest collection of permissively licensed Kotlin code, automatically filtered to include files that have the highest "educational value for learning algorithms in Kotlin".
How to use
python
1from transformers import AutoModelForCausalLM, AutoTokenizer
23# Load pre-trained model and tokenizer4model_name ='JetBrains/CodeLlama-7B-KStack-clean'5tokenizer = AutoTokenizer.from_pretrained(model_name)6model = AutoModelForCausalLM.from_pretrained(model_name).to('cuda')78# Create and encode input9input_text ="""\
10This function takes an integer n and returns factorial of a number:
11fun factorial(n: Int): Int {\
12"""13input_ids = tokenizer.encode(14 input_text, return_tensors='pt'15).to('cuda')1617# Generate18output = model.generate(19 input_ids, max_length=60, num_return_sequences=1,20 pad_token_id=tokenizer.eos_token_id
21)2223# Decode output24generated_text = tokenizer.decode(output[0], skip_special_tokens=True)25print(generated_text)
As with the base model, we can use FIM. To do this, the following format must be used:
'<PRE> ' + prefix + ' <SUF> ' + suffix + ' <MID>'
Training setup
The model was trained on one A100 GPU with following hyperparameters:
Hyperparameter
Value
warmup
100 steps
max_lr
5e-5
scheduler
linear
total_batch_size
32 (~30K tokens per step)
num_epochs
2
More details about fine-tuning can be found in the technical report (coming soon!).
Fine-tuning data
For tuning the model, we used 25K exmaples from the KStack-clean dataset, selected from the larger KStack dataset according to educational value for learning algorithms. In total, the dataset contains about 23M tokens.
Evaluation
For evaluation, we used the Kotlin HumanEval dataset, which contains all 161 tasks from HumanEval translated into Kotlin by human experts. You can find more details about the pre-processing necessary to obtain our results, including the code for running, on the datasets's page.
Here are the results of our evaluation:
Model name
Kotlin HumanEval Pass Rate
CodeLlama-7B
26.89
CodeLlama-7B-KStack-clean
37.89
Ethical Considerations and Limitations
CodeLlama-7B-KStack-clean is a new technology that carries risks with use. The testing conducted to date has not covered, nor could it cover all scenarios. For these reasons, as with all LLMs, CodeLlama-7B-KStack-clean's potential outputs cannot be predicted in advance, and the model may in some instances produce inaccurate or objectionable responses to user prompts. The model was fine-tuned on a specific data format (Kotlin tasks), and deviation from this format can also lead to inaccurate or undesirable responses to user queries. Therefore, before deploying any applications of CodeLlama-7B-KStack-clean, developers should perform safety testing and tuning tailored to their specific applications of the model.
