| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| import torch | |
| tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/DeepSeek-Coder-V2-Lite-Base", trust_remote_code=True) | |
| model = AutoModelForCausalLM.from_pretrained("deepseek-ai/DeepSeek-Coder-V2-Lite-Base", trust_remote_code=True, torch_dtype=torch.bfloat16).cuda() | |
| input_text = """<｜fim▁begin｜>def quick_sort(arr): | |
| if len(arr) <= 1: | |
| return arr | |
| pivot = arr[0] | |
| left = [] | |
| right = [] | |
| <｜fim▁hole｜> | |
| if arr[i] < pivot: | |
| left.append(arr[i]) | |
| else: | |
| right.append(arr[i]) | |
| return quick_sort(left) + [pivot] + quick_sort(right)<｜fim▁end｜>""" | |
| inputs = tokenizer(input_text, return_tensors="pt").to(model.device) | |
| outputs = model.generate(**inputs, max_length=128) | |
| print(tokenizer.decode(outputs[0], skip_special_tokens=True)[len(input_text):]) | |