Model: Add support for HuggingFace config and bad_words_ids

This is necessary for Kobold's API. Current models use bad_words_ids
in generation_config.json, but for some reason, they're also present
in the model's config.json.

Signed-off-by: kingbri <bdashore3@proton.me>

common/transformers_utils.py

@@ -1,6 +1,7 @@
 import json
 import pathlib
 from typing import List, Optional, Union
+from loguru import logger
 from pydantic import BaseModel
 
 
@@ -11,6 +12,7 @@ class GenerationConfig(BaseModel):
     """
 
     eos_token_id: Optional[Union[int, List[int]]] = None
+    bad_words_ids: Optional[List[List[int]]] = None
 
     @classmethod
     def from_file(self, model_directory: pathlib.Path):
@@ -30,3 +32,40 @@ class GenerationConfig(BaseModel):
             return [self.eos_token_id]
         else:
             return self.eos_token_id
+
+
+class HuggingFaceConfig(BaseModel):
+    """
+    An abridged version of HuggingFace's model config.
+    Will be expanded as needed.
+    """
+
+    badwordsids: Optional[str] = None
+
+    @classmethod
+    def from_file(self, model_directory: pathlib.Path):
+        """Create an instance from a generation config file."""
+
+        hf_config_path = model_directory / "config.json"
+        with open(
+            hf_config_path, "r", encoding="utf8"
+        ) as hf_config_json:
+            hf_config_dict = json.load(hf_config_json)
+            return self.model_validate(hf_config_dict)
+
+    def get_badwordsids(self):
+        """Wrapper method to fetch badwordsids."""
+
+        if self.badwordsids:
+            try:
+                bad_words_list = json.loads(self.badwordsids)
+                return bad_words_list
+            except json.JSONDecodeError:
+                logger.warning(
+                    "Skipping badwordsids from config.json "
+                    "since it's not a valid array."
+                )
+
+                return []
+        else:
+            return []

common/utils.py

@@ -18,3 +18,9 @@ def prune_dict(input_dict):
     """Trim out instances of None from a dictionary."""
 
     return {k: v for k, v in input_dict.items() if v is not None}
+
+
+def flat_map(input_list):
+    """Flattens a list of lists into a single list."""
+
+    return [item for sublist in input_list for item in sublist]