Merge branch 'theroyallab:main' into main

Author: gittb

.dockerignore

@@ -0,0 +1,20 @@
+.ruff_cache/
+**/__pycache__/
+venv
+.git
+.gitignore
+.github
+
+# Ignore specific application files
+models/
+loras/
+config.yml
+config_sample.yml
+api_tokens.yml
+api_tokens_sample.yml
+*.bat
+*.sh
+update_scripts
+readme.md
+colab
+start.py

.github/workflows/pages.yml

@@ -48,10 +48,8 @@ jobs:
           npm install @redocly/cli -g
       - name: Export OpenAPI docs
         run: |
-          EXPORT_OPENAPI=1 python main.py
-          mv openapi.json openapi-oai.json
-          EXPORT_OPENAPI=1 python main.py --api-servers kobold
-          mv openapi.json openapi-kobold.json
+          EXPORT_OPENAPI=1 python main.py --openapi-export-path "openapi-oai.json" --api-servers OAI
+          EXPORT_OPENAPI=1 python main.py --openapi-export-path "openapi-kobold.json" --api-servers kobold
       - name: Build and store Redocly site
         run: |
           mkdir static

.gitignore

@@ -214,3 +214,6 @@ openapi.json
 
 # Infinity-emb cache
 .infinity_cache/
+
+# Backup files
+*.bak

backends/exllamav2/grammar.py

@@ -1,9 +1,13 @@
 import traceback
 from exllamav2 import ExLlamaV2, ExLlamaV2Tokenizer
 from exllamav2.generator.filters import ExLlamaV2Filter, ExLlamaV2PrefixFilter
-from lmformatenforcer import JsonSchemaParser, RegexParser
+from lmformatenforcer import (
+    JsonSchemaParser,
+    RegexParser,
+    TokenEnforcer,
+    CharacterLevelParser,
+)
 from lmformatenforcer.integrations.exllamav2 import (
-    ExLlamaV2TokenEnforcerFilter,
     build_token_enforcer_tokenizer_data,
 )
 from loguru import logger
@@ -55,12 +59,48 @@ def feed(self, token):
     def next(self):
         return self.fsm.allowed_token_ids(self.state), set()
 
+    def use_background_worker(self):
+        return True
+
 
 @lru_cache(10)
 def _get_lmfe_tokenizer_data(tokenizer: ExLlamaV2Tokenizer):
     return build_token_enforcer_tokenizer_data(tokenizer)
 
 
+class ExLlamaV2TokenEnforcerFilter(ExLlamaV2Filter):
+    """Filter class for LMFE"""
+
+    token_sequence: List[int]
+
+    def __init__(
+        self,
+        model: ExLlamaV2,
+        tokenizer: ExLlamaV2Tokenizer,
+        character_level_parser: CharacterLevelParser,
+    ):
+        super().__init__(model, tokenizer)
+        tokenizer_data = _get_lmfe_tokenizer_data(tokenizer)
+        self.token_enforcer = TokenEnforcer(tokenizer_data, character_level_parser)
+        self.token_sequence = []
+
+    def begin(self, prefix_str: str):
+        self.token_sequence = []
+
+    def feed(self, token):
+        self.token_sequence.append(int(token[0][0]))
+
+    def next(self):
+        allowed_tokens = self.token_enforcer.get_allowed_tokens(self.token_sequence)
+        if not hasattr(self, "allow_return_type_list"):
+            return set(allowed_tokens), set()
+        else:
+            return sorted(allowed_tokens), []
+
+    def use_background_worker(self):
+        return True
+
+
 def clear_grammar_func_cache():
     """Flush tokenizer_data cache to avoid holding references to
     tokenizers after unloading a model"""
@@ -99,9 +139,7 @@ def add_json_schema_filter(
         # Allow JSON objects or JSON arrays at the top level
         json_prefixes = ["[", "{"]
 
-        lmfilter = ExLlamaV2TokenEnforcerFilter(
-            schema_parser, _get_lmfe_tokenizer_data(tokenizer)
-        )
+        lmfilter = ExLlamaV2TokenEnforcerFilter(model, tokenizer, schema_parser)
         prefix_filter = ExLlamaV2PrefixFilter(model, tokenizer, json_prefixes)
 
         # Append the filters
@@ -110,6 +148,7 @@ def add_json_schema_filter(
     def add_regex_filter(
         self,
         pattern: str,
+        model: ExLlamaV2,
         tokenizer: ExLlamaV2Tokenizer,
     ):
         """Adds an ExllamaV2 filter based on regular expressions."""
@@ -126,9 +165,7 @@ def add_regex_filter(
 
             return
 
-        lmfilter = ExLlamaV2TokenEnforcerFilter(
-            pattern_parser, _get_lmfe_tokenizer_data(tokenizer)
-        )
+        lmfilter = ExLlamaV2TokenEnforcerFilter(model, tokenizer, pattern_parser)
 
         # Append the filters
         self.filters.append(lmfilter)

backends/exllamav2/model.py

@@ -1,5 +1,6 @@
 """The model container class for ExLlamaV2 models."""
 
+import aiofiles
 import asyncio
 import gc
 import math
@@ -16,6 +17,7 @@
     ExLlamaV2Cache_Q4,
     ExLlamaV2Cache_Q6,
     ExLlamaV2Cache_Q8,
+    ExLlamaV2Cache_TP,
     ExLlamaV2Tokenizer,
     ExLlamaV2Lora,
 )
@@ -28,7 +30,7 @@
 from loguru import logger
 from typing import List, Optional, Union
 
-import yaml
+from ruamel.yaml import YAML
 
 from backends.exllamav2.grammar import (
     ExLlamaV2Grammar,
@@ -54,14 +56,6 @@
 from common.transformers_utils import GenerationConfig, HuggingFaceConfig
 from common.utils import coalesce, unwrap
 
-# Dynamic imports
-try:
-    from exllamav2 import ExLlamaV2Cache_TP
-
-    has_tp = True
-except ImportError:
-    has_tp = False
-
 
 class ExllamaV2Container:
     """The model container class for ExLlamaV2 models."""
@@ -106,13 +100,17 @@ class ExllamaV2Container:
     load_lock: asyncio.Lock = asyncio.Lock()
     load_condition: asyncio.Condition = asyncio.Condition()
 
-    def __init__(self, model_directory: pathlib.Path, quiet=False, **kwargs):
+    @classmethod
+    async def create(cls, model_directory: pathlib.Path, quiet=False, **kwargs):
         """
-        Primary initializer for model container.
+        Primary asynchronous initializer for model container.
 
         Kwargs are located in config_sample.yml
         """
 
+        # Create a new instance as a "fake self"
+        self = cls()
+
         self.quiet = quiet
 
         # Initialize config
@@ -155,13 +153,13 @@ def __init__(self, model_directory: pathlib.Path, quiet=False, **kwargs):
             self.draft_config.prepare()
 
         # Create the hf_config
-        self.hf_config = HuggingFaceConfig.from_file(model_directory)
+        self.hf_config = await HuggingFaceConfig.from_file(model_directory)
 
         # Load generation config overrides
         generation_config_path = model_directory / "generation_config.json"
         if generation_config_path.exists():
             try:
-                self.generation_config = GenerationConfig.from_file(
+                self.generation_config = await GenerationConfig.from_file(
                     generation_config_path.parent
                 )
             except Exception:
@@ -171,7 +169,7 @@ def __init__(self, model_directory: pathlib.Path, quiet=False, **kwargs):
                 )
 
         # Apply a model's config overrides while respecting user settings
-        kwargs = self.set_model_overrides(**kwargs)
+        kwargs = await self.set_model_overrides(**kwargs)
 
         # MARK: User configuration
 
@@ -192,17 +190,10 @@ def __init__(self, model_directory: pathlib.Path, quiet=False, **kwargs):
         else:
             # Set tensor parallel
             if use_tp:
-                if has_tp:
-                    self.use_tp = True
+                self.use_tp = True
 
-                    # TP has its own autosplit loader
-                    self.gpu_split_auto = False
-                else:
-                    # TODO: Remove conditional with exl2 v0.1.9 release
-                    logger.warning(
-                        "Tensor parallelism is not supported in the "
-                        "current ExllamaV2 version."
-                    )
+                # TP has its own autosplit loader
+                self.gpu_split_auto = False
 
             # Enable manual GPU split if provided
             if gpu_split:
@@ -320,7 +311,7 @@ def __init__(self, model_directory: pathlib.Path, quiet=False, **kwargs):
             self.cache_size = self.config.max_seq_len
 
         # Try to set prompt template
-        self.prompt_template = self.find_prompt_template(
+        self.prompt_template = await self.find_prompt_template(
             kwargs.get("prompt_template"), model_directory
         )
 
@@ -373,16 +364,25 @@ def __init__(self, model_directory: pathlib.Path, quiet=False, **kwargs):
                 self.draft_config.max_input_len = chunk_size
                 self.draft_config.max_attention_size = chunk_size**2
 
-    def set_model_overrides(self, **kwargs):
+        # Return the created instance
+        return self
+
+    async def set_model_overrides(self, **kwargs):
         """Sets overrides from a model folder's config yaml."""
 
         override_config_path = self.model_dir / "tabby_config.yml"
 
         if not override_config_path.exists():
             return kwargs
 
-        with open(override_config_path, "r", encoding="utf8") as override_config_file:
-            override_args = unwrap(yaml.safe_load(override_config_file), {})
+        async with aiofiles.open(
+            override_config_path, "r", encoding="utf8"
+        ) as override_config_file:
+            contents = await override_config_file.read()
+
+            # Create a temporary YAML parser
+            yaml = YAML(typ="safe")
+            override_args = unwrap(yaml.load(contents), {})
 
             # Merge draft overrides beforehand
             draft_override_args = unwrap(override_args.get("draft"), {})
@@ -393,7 +393,7 @@ def set_model_overrides(self, **kwargs):
             merged_kwargs = {**override_args, **kwargs}
             return merged_kwargs
 
-    def find_prompt_template(self, prompt_template_name, model_directory):
+    async def find_prompt_template(self, prompt_template_name, model_directory):
         """Tries to find a prompt template using various methods."""
 
         logger.info("Attempting to load a prompt template if present.")
@@ -431,7 +431,7 @@ def find_prompt_template(self, prompt_template_name, model_directory):
         # Continue on exception since functions are tried as they fail
         for template_func in find_template_functions:
             try:
-                prompt_template = template_func()
+                prompt_template = await template_func()
                 if prompt_template is not None:
                     return prompt_template
             except TemplateLoadError as e:
@@ -692,7 +692,7 @@ def create_cache(
     ):
         """Utility function to create a model cache."""
 
-        if has_tp and use_tp:
+        if use_tp:
             return ExLlamaV2Cache_TP(
                 model,
                 base=cache_class,
@@ -956,14 +956,6 @@ def check_unsupported_settings(self, **kwargs):
         Meant for dev wheels!
         """
 
-        if unwrap(kwargs.get("dry_allowed_length"), 0) > 0 and not hasattr(
-            ExLlamaV2Sampler.Settings, "dry_multiplier"
-        ):
-            logger.warning(
-                "DRY sampling is not supported by the currently "
-                "installed ExLlamaV2 version."
-            )
-
         return kwargs
 
     async def generate_gen(
@@ -1130,7 +1122,7 @@ async def generate_gen(
         # Add regex filter if it exists
         regex_pattern = unwrap(kwargs.get("regex_pattern"))
         if regex_pattern:
-            grammar_handler.add_regex_filter(regex_pattern, self.tokenizer)
+            grammar_handler.add_regex_filter(regex_pattern, self.model, self.tokenizer)
 
         # Add EBNF filter if it exists
         grammar_string = unwrap(kwargs.get("grammar_string"))

backends/exllamav2/utils.py

@@ -8,7 +8,7 @@
 def check_exllama_version():
     """Verifies the exllama version"""
 
-    required_version = version.parse("0.1.9")
+    required_version = version.parse("0.2.2")
     current_version = version.parse(package_version("exllamav2").split("+")[0])
 
     unsupported_message = (

common/actions.py

@@ -0,0 +1,28 @@
+import json
+from loguru import logger
+
+from common.tabby_config import config, generate_config_file
+from endpoints.server import export_openapi
+from endpoints.utils import do_export_openapi
+
+
+def branch_to_actions() -> bool:
+    """Checks if a optional action needs to be run."""
+
+    if config.actions.export_openapi or do_export_openapi:
+        openapi_json = export_openapi()
+
+        with open(config.actions.openapi_export_path, "w") as f:
+            f.write(json.dumps(openapi_json))
+            logger.info(
+                "Successfully wrote OpenAPI spec to "
+                + f"{config.actions.openapi_export_path}"
+            )
+    elif config.actions.export_config:
+        generate_config_file(filename=config.actions.config_export_path)
+    else:
+        # did not branch
+        return False
+
+    # branched and ran an action
+    return True