Views
No views yet
1# Load the fine-tuned model
2from transformers import AutoTokenizer, AutoModelForCausalLM
3from peft import PeftModel
4import torch
5
6device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
7
8pretrained_checkpoint = 'deepseek-ai/deepseek-coder-1.3b-base'
9finetuned_checkpoint = 'path/to/model'
10
11tokenizer = AutoTokenizer.from_pretrained(finetuned_checkpoint)
12
13old_model = AutoModelForCausalLM.from_pretrained(pretrained_checkpoint)
14old_model.resize_token_embeddings(len(tokenizer))
15
16finetuned_model = PeftModel.from_pretrained(old_model, checkpoint).to(device)
17
18# ----------------------------------------------------------------------------
19# General automatic code completion
20code_example = '''<|secure_function|>\tfunction add('''
21
22model_inputs = tokenizer(code_example, return_tensors="pt").to(device)
23
24input_ids = model_inputs["input_ids"]
25attention_mask = model_inputs["attention_mask"]
26
27generated_ids = finetuned_model.generate(input_ids,
28 do_sample=True,
29 max_length=256,
30 num_beams=4,
31 temperature=0.3,
32 pad_token_id=tokenizer.eos_token_id,
33 attention_mask=attention_mask)
34
35print(tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0])
36
37# Expected output:
38# function add(uint256 a, uint256 b) internal pure returns (uint256) {
39# return a + b;
40# }
41
42# ----------------------------------------------------------------------------
43# Fill-in-the-middle
44def generate_fim(prefix, suffix, model, tokenizer, max_length=256):
45 input_text = f"<|fim_begin|>{prefix}<|fim_hole|>{suffix}<|fim_end|>"
46 inputs = tokenizer.encode(input_text, return_tensors="pt").to(model.device)
47 outputs = model.generate(
48 inputs,
49 max_length=max_length,
50 num_beams=8,
51 temperature=0.3,
52 do_sample=True,
53 pad_token_id=tokenizer.eos_token_id
54 )
55 middle = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)
56 return prefix + middle + suffix
57
58prefix = '''pragma solidity ^0.8.0;\n\n'''
59
60suffix = '''\n\ncontract FOO is Context, IERC20, Ownable {'''
61
62print(generate_fim(prefix, suffix, finetuned_model, tokenizer))
63
64# Expected output:
65# pragma solidity ^0.8.0;
66#
67# import "@openzeppelin/contracts/utils/Context.sol" as Context;
68# import "@openzeppelin/contracts/interfaces/IERC20.sol" as IERC20;
69# import "@openzeppelin/contracts/access/Ownable.sol" as Ownable;
70#
71# contract FOO is Context, IERC20, Ownable {
721@misc{hensel2025fim_model,
2 title = {Finetuned deepseek-coder-1.3b-base model for automatic code completion of Solidity code},
3 author={Fabian Hensel},
4 year={2025}
5}