Views
No views yet
1LANGS = [
2 "Python",
3 "Rust",
4 "JavaScript",
5 "Java",
6 "Go",
7 "C++",
8 "C#",
9 "Ruby",
10 "PHP",
11 "TypeScript",
12 "C",
13 "Scala",
14 "Swift",
15 "Kotlin",
16 "Objective-C",
17 "Perl",
18 "Haskell",
19 "Bash",
20 "Sh",
21 "Lua",
22 "R",
23 "Julia",
24]1import difflib
2class NDiff:
3 def __init__(self, s1, s2):
4 self.s1 = s1
5 self.s2 = s2
6 self.diff = difflib.ndiff(s1.split("\n"), s2.split("\n"))
7
8 def __str__(self):
9 return "\n".join([l for l in self.diff if l[0] != "?"])
10
11 def str_colored(self):
12 import colored
13
14 buf = ""
15 for l in self.diff:
16 if l[0] == "?":
17 continue
18 if l[0] == "-":
19 buf += colored.stylize(l, colored.fg("red"))
20 elif l[0] == "+":
21 buf += colored.stylize(l, colored.fg("green"))
22 else:
23 buf += l
24 buf += "\n"
25 return buf
26
27 def num_removed(self):
28 return len([l for l in self.diff if l[0] == "-"])
29
30 def num_added(self):
31 return len([l for l in self.diff if l[0] == "+"])
32
33 def __repr__(self):
34 return self.__str__()
35
36def format_prompt(old, new):
37 diff_header = "<diff>"
38 instr_header = "<commit_message>"
39 diff = str(NDiff(old, new))
40 return f"{diff_header}\n{diff}\n{instr_header}\n"
41
42def gen(old, new, max_new_tokens=200, temperature=0.45, top_p=0.90):
43 prompt = format_prompt(old, new)
44 toks = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
45 outs = model.generate(toks, max_new_tokens=max_new_tokens, do_sample=True, temperature=temperature, top_p=top_p)
46 return [tokenizer.decode(out[len(toks[0]):], skip_special_tokens=True) for out in outs]1- import datasets
2- from pathlib import Path
3 from code_editing.models import CodeLlamaEditModel, LlamaChatModel, EditModel, EditCommand, ChatAdaptorEditModel, OctoCoderChatModel, codellama_edit_prompt_diff, apply_rel_diff_trim, OpenAIChatModel, StarCoderCommitEditModel
4 from code_editing.humanevalpack import batch_prompts_from_example
5 from code_editing.utils import gunzip_json_write
6 from typing import List, Callable
7 from tqdm import tqdm
8
9
10 # NOTE: this is the factory for each model type. to add a new model type, add a new case here
11 # and implement it in models.py. Also, add a new case in the argument parser below.
12- def model_factory(model_type: str, quantize=False, num_gpus=1) -> Callable[[str], EditModel]:
13+ def model_factory(
14+ model_type: str,
15+ quantize=False,
16+ num_gpus=1,
17+ system_supported=True,
18+ ) -> Callable[[str], EditModel]:
19 if model_type == "codellama" or model_type == "deepseek":
20 return CodeLlamaEditModel
21 elif model_type == "starcoder":
22 return StarCoderCommitEditModel
23 elif model_type == "codellama-diff":
24 return (lambda path: CodeLlamaEditModel(path, prompt_format=codellama_edit_prompt_diff, post_process=apply_rel_diff_trim))
25 elif model_type == "openai":
26 return (lambda path: ChatAdaptorEditModel(OpenAIChatModel(path)))
27 elif model_type == "codellama-chat":
28- return (lambda path: ChatAdaptorEditModel(LlamaChatModel(path, quantization=quantize, num_gpus=num_gpus)))
29+ return (lambda path: ChatAdaptorEditModel(LlamaChatModel(path, quantization=quantize, num_gpus=num_gpus, system_supported=system_supported)))
30 elif model_type == "octocoder":
31 return (lambda path: ChatAdaptorEditModel(OctoCoderChatModel(path, quantization=quantize, num_gpus=num_gpus)))
32 else:
33 raise ValueError(f"Unknown model type: {model_type}")
34
35 def complete_problem(example: EditCommand, model: EditModel, batch_size: int, completion_limit: int, **kwargs) -> List[str]:
36 batches = batch_prompts_from_example(example, batch_size, completion_limit)
37
38 completions = []
39 for batch in batches:
40 resps = model.generate(batch, **kwargs)
41 for resp in resps:
42 completions.append(resp["content"])
43
44 return completionsAdd system_supported argument to model_factory