DeciCoder is a compact language model for generating Python, JavaScript, and Java. Its 1B-parameter size makes it attractive for a coding copilot that can run locally or potentially fit inside an extension. section 10.1 section 10.5
Its specialization is also its main trade-off. Training on a filtered Stack dataset limited to three languages helps code generation, while the 2,048-token context window restricts how much retrieved code a RAG prompt can include. section 10.1
Use DeciCoder with these production implications in mind:
- Expect strong performance when half-written code is the input and generated code is the output. section 10.4
- Treat instruction-style prompts and few-shot examples cautiously because the model was not instruction tuned and may generalize poorly to unfamiliar task formats. section 10.1 section 10.4
- Improve results by aligning prompts with the model’s training format; plain-English instructions may be less useful than relevant code examples. section 10.4 section 10.5
Local coding-copilot service
A VS Code extension can call a locally served DeciCoder API. A minimal request sends a prompt such as `def hello_world(name):` to a FastAPI endpoint, which returns generated code. section 10.1
Listing 10.1 A simple FastAPI endpoint using DeciCoder
import argparse
from fastapi import FastAPI, Request
from fastapi.responses import Response
import torch
import uvicorn
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
StoppingCriteria,
StoppingCriteriaList,
)
torch.backends.cuda.enable_mem_efficient_sdp(False) #1
torch.backends.cuda.enable_flash_sdp(False)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
stop_tokens = ["def", "class", "Instruction", "Output"] #2
stop_token_ids = [589, 823, 9597, 2301]
class StopOnTokens(StoppingCriteria):
def __call__(
self,
input_ids: torch.LongTensor,
scores: torch.FloatTensor,
**kwargs,
) -> bool:
stop_ids = stop_token_ids
for stop_id in stop_ids:
if input_ids[0][-1] == stop_id:
return True
return False
tokenizer = AutoTokenizer.from_pretrained("Deci/DeciCoder-1b") #3
tokenizer.add_special_tokens( #3
{"additional_special_tokens": stop_tokens}, #3
replace_additional_special_tokens=False, #3
) #3
model = AutoModelForCausalLM.from_pretrained( #3
"Deci/DeciCoder-1b", torch_dtype=torch.bfloat16, trust_remote_code=True #3
) #3
model = model.to(device) #3
app = FastAPI() #4
@app.post("/generate")
async def generate(request: Request) -> Response:
"""Generate LLM Response
The request should be a JSON object with the following fields:
- prompt: the prompt to use for the generation.
"""
request_dict = await request.json()
prompt = request_dict.pop("prompt")
# ... #5
inputs = tokenizer(prompt, return_tensors="pt").to(device) #6
response_tokens = model.generate(
inputs["input_ids"],
max_new_tokens=1024,
stopping_criteria=StoppingCriteriaList([StopOnTokens()]),
)
input_length = inputs["input_ids"].shape[1]
response = tokenizer.decode(
response_tokens[0][input_length:], skip_special_tokens=True
)
return response
if __name__ == "__main__":
parser = argparse.ArgumentParser() #7
parser.add_argument("--host", type=str, default=None)
parser.add_argument("--port", type=int, default=8000)
args = parser.parse_args()
uvicorn.run(app, host=args.host, port=args.port, log_level="debug")
Inspect the FastAPI implementation, model loading, device selection, and stopping criteria used to serve DeciCoder. section 10.1
