Local llama support

aider/__init__.py

@@ -1 +1 @@
-__version__ = "0.7.2"
+__version__ = "0.7.3"

aider/coders/base_coder.py

@@ -23,6 +23,7 @@
 from aider.repomap import RepoMap
 
 from ..dump import dump  # noqa: F401
+from ..tokenizers import OPENAI
 
 
 class MissingAPIKeyError(ValueError):
@@ -114,6 +115,7 @@ def __init__(
         code_theme="default",
         stream=True,
         use_git=True,
+        tokenizer=OPENAI,
     ):
         if not fnames:
             fnames = []
@@ -152,6 +154,8 @@ def __init__(
 
         self.show_diffs = show_diffs
 
+        self.tokenizer = tokenizer
+
         self.commands = Commands(self.io, self)
 
         if use_git:

aider/commands.py

@@ -6,16 +6,25 @@
 
 import git
 import tiktoken
+import transformers
 from prompt_toolkit.completion import Completion
+from transformers import LlamaTokenizer
 
 from aider import prompts, utils
+from aider.tokenizers import OPENAI, LLAMA
 
 
 class Commands:
     def __init__(self, io, coder):
         self.io = io
         self.coder = coder
-        self.tokenizer = tiktoken.encoding_for_model(coder.main_model.name)
+
+        if coder.tokenizer == OPENAI:
+            self.tokenizer = tiktoken.encoding_for_model(coder.main_model.name)
+        elif coder.tokenizer == LLAMA:
+            self.tokenizer = transformers.LlamaTokenizer
+        else:
+            raise ValueError(f"No such tokenizer found for {coder.tokenizer}")
 
     def is_command(self, inp):
         if inp[0] == "/":

aider/main.py

@@ -7,6 +7,7 @@
 from aider import __version__, models
 from aider.coders import Coder
 from aider.io import InputOutput
+from aider.tokenizers import OPENAI
 
 
 def get_git_root():
@@ -93,6 +94,21 @@ def main(args=None, input=None, output=None):
         const=models.GPT35_16k.name,
         help=f"Use {models.GPT35_16k.name} model for the main chat (gpt-4 is better)",
     )
+    parser.add_argument(
+        "--model-tokens",
+        metavar="MODEL_TOKENS",
+        dest="model_tokens",
+        default=None,
+        help=f"Specify the number of tokens to use with the model (only applicable when using local models). This "
+             f"value will be multiplied by 1024 to determine the actual token count.",
+    )
+    parser.add_argument(
+        "--tokenizer",
+        metavar="TOKENIZER",
+        dest="tokenizer",
+        default=OPENAI,
+        help=f"Specify the the number of tokens to use with the model (only applicable when using local models).",
+    )
     parser.add_argument(
         "--edit-format",
         metavar="EDIT_FORMAT",
@@ -259,7 +275,13 @@ def main(args=None, input=None, output=None):
         io.tool_error("No OpenAI API key provided. Use --openai-api-key or env OPENAI_API_KEY.")
         return 1
 
-    main_model = models.Model(args.model)
+    try:
+        main_model = models.get_model(args.model)
+    except:
+        if not args.model_tokens:
+            io.tool_error("Model tokens must be specified if using a local model")
+            return 1
+        main_model = models.Model(args.model, 2, "whole")
 
     coder = Coder.create(
         main_model,
@@ -280,6 +302,7 @@ def main(args=None, input=None, output=None):
         code_theme=args.code_theme,
         stream=args.stream,
         use_git=args.git,
+        tokenizer=args.tokenizer,
     )
 
     if args.dirty_commits:

aider/models.py

@@ -1,11 +1,3 @@
-import re
-
-known_tokens = {
-    "gpt-3.5-turbo": 4,
-    "gpt-4": 8,
-}
-
-
 class Model:
     always_available = False
     use_repo_map = False
@@ -14,26 +6,16 @@ class Model:
     prompt_price = None
     completion_price = None
 
-    def __init__(self, name):
+    def __init__(self, name, tokens, edit_format):
         self.name = name
-
-        tokens = None
-
-        match = re.search(r"-([0-9]+)k", name)
-        if match:
-            tokens = int(match.group(1))
-        else:
-            for m, t in known_tokens.items():
-                if name.startswith(m):
-                    tokens = t
+        self.edit_format = edit_format
 
         if tokens is None:
             raise ValueError(f"Unknown context window size for model: {name}")
 
         self.max_context_tokens = tokens * 1024
 
         if self.is_gpt4():
-            self.edit_format = "diff"
             self.use_repo_map = True
             self.send_undo_reply = True
 
@@ -47,7 +29,6 @@ def __init__(self, name):
             return
 
         if self.is_gpt35():
-            self.edit_format = "whole"
             self.always_available = True
 
             if tokens == 4:
@@ -59,7 +40,7 @@ def __init__(self, name):
 
             return
 
-        raise ValueError(f"Unsupported model: {name}")
+        return
 
     def is_gpt4(self):
         return self.name.startswith("gpt-4")
@@ -71,6 +52,17 @@ def __str__(self):
         return self.name
 
 
-GPT4 = Model("gpt-4")
-GPT35 = Model("gpt-3.5-turbo")
-GPT35_16k = Model("gpt-3.5-turbo-16k")
+GPT4 = Model("gpt-4", 8, "diff")
+GPT35 = Model("gpt-3.5-turbo", 4, "whole")
+GPT35_16k = Model("gpt-3.5-turbo-16k", 4, "whole")
+
+
+def get_model(name):
+    if name == GPT4.name:
+        return GPT4
+    elif name == GPT35.name:
+        return GPT35
+    elif name == GPT35_16k:
+        return GPT35_16k
+
+    raise ValueError(f"No such model: {name}")

aider/tokenizers.py

@@ -0,0 +1,2 @@
+OPENAI = "openai"
+LLAMA = "llama"

requirements.txt

@@ -23,11 +23,15 @@ wcwidth==0.2.6
 yarl==1.9.2
 pytest==7.3.1
 tiktoken==0.4.0
-configargparse
+configargparse~=1.5.5
 PyYAML
 backoff==2.2.1
 networkx==3.1
 diskcache==5.6.1
 numpy==1.24.3
 scipy==1.10.1
 jsonschema==4.17.3
+transformers==4.30.2
+
+packaging~=23.1
+setuptools~=58.1.0